<commit_before>package beater

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"os"
	"time"

	"github.com/elastic/beats/libbeat/beat"
	"github.com/elastic/beats/libbeat/common"
	"github.com/elastic/beats/libbeat/logp"
	"github.com/elastic/beats/libbeat/publisher"
	"github.com/joshuar/pingbeat/config"
	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
	"gopkg.in/go-playground/pool.v3"
)

const pingTimeout = 4 * time.Second

// Pingbeat contains configuration details
type Pingbeat struct {
	done        chan struct{}
	config      config.Config
	client      publisher.Client
	ipv4network string
	ipv6network string
	targets     map[string]Target
}

// PingInfo contains details about active ping requests/replies
type PingInfo struct {
	ID         int
	Seq        int
	Target     string
	Sent       time.Time
	Received   time.Time
	RTT        time.Duration
	Loss       bool
	LossReason string
}

// New creates a new Pingbeat beater struct
func New(b *beat.Beat, cfg *common.Config) (beat.Beater, error) {
	config := config.DefaultConfig
	if err := cfg.Unpack(&config); err != nil {
		return nil, fmt.Errorf("Error reading config file: %v", err)
	}

	bt := &Pingbeat{
		done:   make(chan struct{}),
		config: config,
	}

	// Use privileged (i.e. raw socket) ping by default, else use a UDP ping
	if bt.config.Privileged {
		if os.Getuid() != 0 {
			return nil, fmt.Errorf("privileged specified but not running with privileges")
		}
		bt.ipv4network = "ip4:icmp"
		bt.ipv6network = "ip6:ipv6-icmp"
	} else {
		bt.ipv4network = "udp4"
		bt.ipv6network = "udp6"
	}

	// Fill the IPv4/IPv6 targets maps
	bt.targets = NewTargets(bt.config.Targets, bt.config.Privileged, bt.config.UseIPv4, bt.config.UseIPv6)
	return bt, nil
}

// Run is the main loop which sends/recieves ICMP messages and cleans up stale
// requests
func (bt *Pingbeat) Run(b *beat.Beat) error {
	logp.Info("pingbeat is running! Hit CTRL-C to stop it.")

	bt.client = b.Publisher.Connect()

	// Set up send/receive pools
	spool := pool.NewLimited(uint(len(bt.targets)) * uint(pingTimeout.Seconds()))
	defer spool.Close()

	// Set up a ticker to loop for the period specified
	ticker := time.NewTicker(bt.config.Period)
	defer ticker.Stop()
	timeout := time.NewTicker(pingTimeout)
	defer timeout.Stop()

	// Create a new global state to track active ping requests
	state := NewPingState()

	// Start receivers to capture incoming ping replies
	// Create required connections
	var ipv4conn, ipv6conn *icmp.PacketConn
	var err error
	var pingID = os.Getpid() & 0xffff
	logp.Debug("pingbeat", "pingID: %v", pingID)
	if bt.config.UseIPv4 {
		if ipv4conn, err = createConn(bt.ipv4network, "0.0.0.0"); err != nil {
			logp.Err("Error creating %s connection: %v", bt.ipv4network, err)
			return nil
		}
		logp.Info("Using %s connection", bt.ipv4network)
		go RecvPings(pingID, bt, state, ipv4conn)
	}
	if bt.config.UseIPv6 {
		if ipv6conn, err = createConn(bt.ipv6network, "::"); err != nil {
			logp.Err("Error creating %s connection: %v", bt.ipv6network, err)
			return nil
		}
		logp.Info("Using %s connection", bt.ipv6network)
		go RecvPings(pingID, bt, state, ipv6conn)
	}

	for {
		select {
		case <-bt.done:
			return nil
		case <-timeout.C:
			// Timeout reached, clean up any pending ping requests where there
			// has been no response
			go state.CleanPings(pingTimeout)
		case <-ticker.C:
			// Batch queue echo request
			sendBatch := spool.Batch()
			go func(*icmp.PacketConn, *icmp.PacketConn) {
				for ip, target := range bt.targets {
					if net.ParseIP(ip).To4() != nil {
						sendBatch.Queue(SendPing(ipv4conn, pingTimeout, state.GetSeqNo(), target.Addr))
					} else {
						sendBatch.Queue(SendPing(ipv6conn, pingTimeout, state.GetSeqNo(), target.Addr))
					}
				}
				sendBatch.QueueComplete()
			}(ipv4conn, ipv6conn)

			// For each successfully sent echo request
			for result := range sendBatch.Results() {
				// Grab info of the sent request
				if result.Value() == nil {
					logp.Debug("pingbeat", "Send unsuccessful: %v", result.Error())
					break
				}
				info := result.Value().(*PingInfo)
				if err := result.Error(); err != nil {
					logp.Debug("pingbeat", "Send unsuccessful: %v", err)
				}
				success := state.AddPing(info.Target, info.Seq, info.Sent)
				if !success {
					logp.Err("Error adding ping (%v:%v) to state", info.Seq, info.Target)
				}
			}
		}
	}
}

// Stop cleans up Pingbeat
func (bt *Pingbeat) Stop() {
	bt.client.Close()
	close(bt.done)
}

// RecvPings listens for ICMP messages, decodes them into the right type and
// checks if they were sent by this Pingbeat, before processing them
func RecvPings(myID int, bt *Pingbeat, state *PingState, conn *icmp.PacketConn) {
	for {
		// Based on the connection, work out whether we are dealing with
		// IPv4 or IPv6 ICMP messages
		var pingType icmp.Type
		switch {
		case conn.IPv4PacketConn() != nil:
			pingType = ipv4.ICMPTypeEcho
		case conn.IPv4PacketConn() != nil:
			pingType = ipv6.ICMPTypeEchoRequest
		default:
			err := errors.New("Unknown connection type")
			logp.Err("Error parsing connection: %v", err)
			break
		}

		// Read data from the connection
		bd := make([]byte, 1500)
		n, peer, err := conn.ReadFrom(bd)
		if err != nil {
			logp.Err("Couldn't read from connection: %v", err)
			continue
		}
		var target string
		switch peer.(type) {
		case *net.UDPAddr:
			target, _, _ = net.SplitHostPort(peer.String())
		case *net.IPAddr:
			target = peer.String()
		default:
			logp.Err("Error parsing received address %v", target)
			continue
		}

		if n == 0 {
			continue
		}
		// Parse the data into an ICMP message
		message, err := icmp.ParseMessage(pingType.Protocol(), bd[:n])
		if err != nil {
			logp.Err("Couldn't parse response: %v", err)
			continue
		}

		ping := &PingInfo{}
		// Switch for the ICMP message type
		switch message.Body.(type) {
		case *icmp.Echo:
			ping.Seq = message.Body.(*icmp.Echo).Seq
			ping.ID = message.Body.(*icmp.Echo).ID
			ping.Target = target
			ping.Loss = false
			ping.Received = time.Now().UTC()
		case *icmp.TimeExceeded:
			ping.Loss = true
			ping.LossReason = "Time Exceeded"
			ping.ID, ping.Seq, ping.Target = parseICMPError(message.Body.(*icmp.TimeExceeded).Data)
		case *icmp.PacketTooBig:
			ping.Loss = true
			ping.LossReason = "Packet Too Big"
			ping.ID, ping.Seq, ping.Target = parseICMPError(message.Body.(*icmp.PacketTooBig).Data)
		case *icmp.DstUnreach:
			ping.Loss = true
			ping.LossReason = "Destination Unreachable"
			ping.ID, ping.Seq, ping.Target = parseICMPError(message.Body.(*icmp.DstUnreach).Data)
		default:
		}
		if ping.ID != 0 && ping.ID != myID {
			logp.Debug("RecvPings", "Ping response from %v not from me:", target)
		} else {
			if !ping.Loss {
				ping.RTT = state.CalcPingRTT(ping.Seq, ping.Received)
			} else {
				logp.Warn("%v: %v", ping.LossReason, ping.Target)
			}
			go bt.ProcessPing(ping)
			state.DelPing(ping.Seq)
		}
	}
}

// SendPing sends an ICMP EchoRequest packet to with provided sequence number to
// the provided target through the given connection
func SendPing(conn *icmp.PacketConn, timeout time.Duration, seq int, addr net.Addr) pool.WorkFunc {
	return func(wu pool.WorkUnit) (interface{}, error) {
		if wu.IsCancelled() {
			logp.Debug("SendPings", "SendPing: workunit cancelled")
			return nil, nil
		}
		// Based on the connection, work out whether we are dealing with
		// IPv4 or IPv6 ICMP messages
		var pingType icmp.Type
		switch {
		case conn.IPv4PacketConn() != nil:
			pingType = ipv4.ICMPTypeEcho
		case conn.IPv4PacketConn() != nil:
			pingType = ipv6.ICMPTypeEchoRequest
		default:
			err := errors.New("Unknown connection type")
			return nil, err
		}

		// Create an ICMP Echo Request
		var id = os.Getpid() & 0xffff
		message := &icmp.Message{
			Type: pingType, Code: 0,
			Body: &icmp.Echo{
				ID:   id,
				Seq:  seq,
				Data: []byte("pingbeat: y'know, for pings!"),
			},
		}
		// Marshall the Echo request for sending via a connection
		binary, err := message.Marshal(nil)
		if err != nil {
			return nil, err
		}
		var t string
		switch addr.(type) {
		case *net.UDPAddr:
			t, _, _ = net.SplitHostPort(addr.String())
		case *net.IPAddr:
			t = addr.String()
		default:
			err := errors.New("Unknown address type")
			return nil, err
		}

		ping := &PingInfo{
			Seq:    seq,
			Target: t,
		}
		// Send the request
		if _, err := conn.WriteTo(binary, addr); err != nil {
			return ping, err
		}
		ping.Sent = time.Now().UTC()
		return ping, nil
	}
}

// ProcessPing fetches the details of this ping from the current state
// and then creates an ping event to be published
func (bt *Pingbeat) ProcessPing(ping *PingInfo) {
	if _, found := bt.targets[ping.Target]; !found {
		logp.Err("No details for %v in targets!", ping.Target)
	} else {
		name := bt.targets[ping.Target].Name
		tags := bt.targets[ping.Target].Tags
		if ping.Loss {
			event := common.MapStr{
				"@timestamp":  common.Time(time.Now().UTC()),
				"type":        "pingbeat",
				"target.name": name,
				"target.addr": ping.Target,
				"target.tags": tags,
				"loss":        true,
				"reason":      ping.LossReason,
			}
			go bt.client.PublishEvent(event)
			logp.Debug("ProcessPing", "Processed ping error for %v (%v): %v", name, ping.Target, ping.LossReason)
		} else {
			event := common.MapStr{
				"@timestamp":  common.Time(time.Now().UTC()),
				"type":        "pingbeat",
				"target.name": name,
				"target.addr": ping.Target,
				"target.tags": tags,
				"rtt":         milliSeconds(ping.RTT),
			}
			go bt.client.PublishEvent(event)
			logp.Debug("ProcessPing", "Processed ping %v for %v (%v): %v", ping.Seq, name, ping.Target, ping.RTT)
		}
	}
}

func parseICMPError(data []byte) (int, int, string) {
	IPheader, err := ipv4.ParseHeader(data[:len(data)-8])
	if err != nil {
		logp.Err("parseICMPError", "Failed to parse packet header:", err)
	}
	ICMPHdr := data[IPheader.Len:]
	var ID, Seq uint16
	err = binary.Read(bytes.NewReader(ICMPHdr[6:8]), binary.BigEndian, &Seq)
	if err != nil {
		logp.Err("parseICMPError", "Failed to parse packet header:", err)
	}
	err = binary.Read(bytes.NewReader(ICMPHdr[4:6]), binary.BigEndian, &ID)
	if err != nil {
		logp.Err("parseICMPError", "Failed to parse packet header:", err)
	}
	return int(ID), int(Seq), IPheader.Dst.String()
}

func createConn(n string, a string) (*icmp.PacketConn, error) {
	c, err := icmp.ListenPacket(n, a)
	if err != nil {
		return nil, err
	}
	return c, nil
}

func milliSeconds(d time.Duration) float64 {
	msec := d / time.Millisecond
	nsec := d % time.Millisecond
	return float64(msec) + float64(nsec)*1e-6
}
<commit_msg>Adjust published fields.<commit_after>package beater

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"os"
	"time"

	"github.com/elastic/beats/libbeat/beat"
	"github.com/elastic/beats/libbeat/common"
	"github.com/elastic/beats/libbeat/logp"
	"github.com/elastic/beats/libbeat/publisher"
	"github.com/joshuar/pingbeat/config"
	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
	"gopkg.in/go-playground/pool.v3"
)

const pingTimeout = 4 * time.Second

// Pingbeat contains configuration details
type Pingbeat struct {
	done        chan struct{}
	config      config.Config
	client      publisher.Client
	ipv4network string
	ipv6network string
	targets     map[string]Target
}

// PingInfo contains details about active ping requests/replies
type PingInfo struct {
	ID         int
	Seq        int
	Target     string
	Sent       time.Time
	Received   time.Time
	RTT        time.Duration
	Loss       bool
	LossReason string
}

// New creates a new Pingbeat beater struct
func New(b *beat.Beat, cfg *common.Config) (beat.Beater, error) {
	config := config.DefaultConfig
	if err := cfg.Unpack(&config); err != nil {
		return nil, fmt.Errorf("Error reading config file: %v", err)
	}

	bt := &Pingbeat{
		done:   make(chan struct{}),
		config: config,
	}

	// Use privileged (i.e. raw socket) ping by default, else use a UDP ping
	if bt.config.Privileged {
		if os.Getuid() != 0 {
			return nil, fmt.Errorf("privileged specified but not running with privileges")
		}
		bt.ipv4network = "ip4:icmp"
		bt.ipv6network = "ip6:ipv6-icmp"
	} else {
		bt.ipv4network = "udp4"
		bt.ipv6network = "udp6"
	}

	// Fill the IPv4/IPv6 targets maps
	bt.targets = NewTargets(bt.config.Targets, bt.config.Privileged, bt.config.UseIPv4, bt.config.UseIPv6)
	return bt, nil
}

// Run is the main loop which sends/recieves ICMP messages and cleans up stale
// requests
func (bt *Pingbeat) Run(b *beat.Beat) error {
	logp.Info("pingbeat is running! Hit CTRL-C to stop it.")

	bt.client = b.Publisher.Connect()

	// Set up send/receive pools
	spool := pool.NewLimited(uint(len(bt.targets)) * uint(pingTimeout.Seconds()))
	defer spool.Close()

	// Set up a ticker to loop for the period specified
	ticker := time.NewTicker(bt.config.Period)
	defer ticker.Stop()
	timeout := time.NewTicker(pingTimeout)
	defer timeout.Stop()

	// Create a new global state to track active ping requests
	state := NewPingState()

	// Start receivers to capture incoming ping replies
	// Create required connections
	var ipv4conn, ipv6conn *icmp.PacketConn
	var err error
	var pingID = os.Getpid() & 0xffff
	logp.Debug("pingbeat", "pingID: %v", pingID)
	if bt.config.UseIPv4 {
		if ipv4conn, err = createConn(bt.ipv4network, "0.0.0.0"); err != nil {
			logp.Err("Error creating %s connection: %v", bt.ipv4network, err)
			return nil
		}
		logp.Info("Using %s connection", bt.ipv4network)
		go RecvPings(pingID, bt, state, ipv4conn)
	}
	if bt.config.UseIPv6 {
		if ipv6conn, err = createConn(bt.ipv6network, "::"); err != nil {
			logp.Err("Error creating %s connection: %v", bt.ipv6network, err)
			return nil
		}
		logp.Info("Using %s connection", bt.ipv6network)
		go RecvPings(pingID, bt, state, ipv6conn)
	}

	for {
		select {
		case <-bt.done:
			return nil
		case <-timeout.C:
			// Timeout reached, clean up any pending ping requests where there
			// has been no response
			go state.CleanPings(pingTimeout)
		case <-ticker.C:
			// Batch queue echo request
			sendBatch := spool.Batch()
			go func(*icmp.PacketConn, *icmp.PacketConn) {
				for ip, target := range bt.targets {
					if net.ParseIP(ip).To4() != nil {
						sendBatch.Queue(SendPing(ipv4conn, pingTimeout, state.GetSeqNo(), target.Addr))
					} else {
						sendBatch.Queue(SendPing(ipv6conn, pingTimeout, state.GetSeqNo(), target.Addr))
					}
				}
				sendBatch.QueueComplete()
			}(ipv4conn, ipv6conn)

			// For each successfully sent echo request
			for result := range sendBatch.Results() {
				// Grab info of the sent request
				if result.Value() == nil {
					logp.Debug("pingbeat", "Send unsuccessful: %v", result.Error())
					break
				}
				info := result.Value().(*PingInfo)
				if err := result.Error(); err != nil {
					logp.Debug("pingbeat", "Send unsuccessful: %v", err)
				}
				success := state.AddPing(info.Target, info.Seq, info.Sent)
				if !success {
					logp.Err("Error adding ping (%v:%v) to state", info.Seq, info.Target)
				}
			}
		}
	}
}

// Stop cleans up Pingbeat
func (bt *Pingbeat) Stop() {
	bt.client.Close()
	close(bt.done)
}

// RecvPings listens for ICMP messages, decodes them into the right type and
// checks if they were sent by this Pingbeat, before processing them
func RecvPings(myID int, bt *Pingbeat, state *PingState, conn *icmp.PacketConn) {
	for {
		// Based on the connection, work out whether we are dealing with
		// IPv4 or IPv6 ICMP messages
		var pingType icmp.Type
		switch {
		case conn.IPv4PacketConn() != nil:
			pingType = ipv4.ICMPTypeEcho
		case conn.IPv4PacketConn() != nil:
			pingType = ipv6.ICMPTypeEchoRequest
		default:
			err := errors.New("Unknown connection type")
			logp.Err("Error parsing connection: %v", err)
			break
		}

		// Read data from the connection
		bd := make([]byte, 1500)
		n, peer, err := conn.ReadFrom(bd)
		if err != nil {
			logp.Err("Couldn't read from connection: %v", err)
			continue
		}
		var target string
		switch peer.(type) {
		case *net.UDPAddr:
			target, _, _ = net.SplitHostPort(peer.String())
		case *net.IPAddr:
			target = peer.String()
		default:
			logp.Err("Error parsing received address %v", target)
			continue
		}

		if n == 0 {
			continue
		}
		// Parse the data into an ICMP message
		message, err := icmp.ParseMessage(pingType.Protocol(), bd[:n])
		if err != nil {
			logp.Err("Couldn't parse response: %v", err)
			continue
		}

		ping := &PingInfo{}
		// Switch for the ICMP message type
		switch message.Body.(type) {
		case *icmp.Echo:
			ping.Seq = message.Body.(*icmp.Echo).Seq
			ping.ID = message.Body.(*icmp.Echo).ID
			ping.Target = target
			ping.Loss = false
			ping.Received = time.Now().UTC()
		case *icmp.TimeExceeded:
			ping.Loss = true
			ping.LossReason = "Time Exceeded"
			ping.ID, ping.Seq, ping.Target = parseICMPError(message.Body.(*icmp.TimeExceeded).Data)
		case *icmp.PacketTooBig:
			ping.Loss = true
			ping.LossReason = "Packet Too Big"
			ping.ID, ping.Seq, ping.Target = parseICMPError(message.Body.(*icmp.PacketTooBig).Data)
		case *icmp.DstUnreach:
			ping.Loss = true
			ping.LossReason = "Destination Unreachable"
			ping.ID, ping.Seq, ping.Target = parseICMPError(message.Body.(*icmp.DstUnreach).Data)
		default:
		}
		if ping.ID != 0 && ping.ID != myID {
			logp.Debug("RecvPings", "Ping response from %v not from me:", target)
		} else {
			if !ping.Loss {
				ping.RTT = state.CalcPingRTT(ping.Seq, ping.Received)
			} else {
				logp.Warn("%v: %v", ping.LossReason, ping.Target)
			}
			go bt.ProcessPing(ping)
			state.DelPing(ping.Seq)
		}
	}
}

// SendPing sends an ICMP EchoRequest packet to with provided sequence number to
// the provided target through the given connection
func SendPing(conn *icmp.PacketConn, timeout time.Duration, seq int, addr net.Addr) pool.WorkFunc {
	return func(wu pool.WorkUnit) (interface{}, error) {
		if wu.IsCancelled() {
			logp.Debug("SendPings", "SendPing: workunit cancelled")
			return nil, nil
		}
		// Based on the connection, work out whether we are dealing with
		// IPv4 or IPv6 ICMP messages
		var pingType icmp.Type
		switch {
		case conn.IPv4PacketConn() != nil:
			pingType = ipv4.ICMPTypeEcho
		case conn.IPv4PacketConn() != nil:
			pingType = ipv6.ICMPTypeEchoRequest
		default:
			err := errors.New("Unknown connection type")
			return nil, err
		}

		// Create an ICMP Echo Request
		var id = os.Getpid() & 0xffff
		message := &icmp.Message{
			Type: pingType, Code: 0,
			Body: &icmp.Echo{
				ID:   id,
				Seq:  seq,
				Data: []byte("pingbeat: y'know, for pings!"),
			},
		}
		// Marshall the Echo request for sending via a connection
		binary, err := message.Marshal(nil)
		if err != nil {
			return nil, err
		}
		var t string
		switch addr.(type) {
		case *net.UDPAddr:
			t, _, _ = net.SplitHostPort(addr.String())
		case *net.IPAddr:
			t = addr.String()
		default:
			err := errors.New("Unknown address type")
			return nil, err
		}

		ping := &PingInfo{
			Seq:    seq,
			Target: t,
		}
		// Send the request
		if _, err := conn.WriteTo(binary, addr); err != nil {
			return ping, err
		}
		ping.Sent = time.Now().UTC()
		return ping, nil
	}
}

// ProcessPing fetches the details of this ping from the current state
// and then creates an ping event to be published
func (bt *Pingbeat) ProcessPing(ping *PingInfo) {
	if _, found := bt.targets[ping.Target]; !found {
		logp.Err("No details for %v in targets!", ping.Target)
	} else {
		name := bt.targets[ping.Target].Name
		tags := bt.targets[ping.Target].Tags
		if ping.Loss {
			event := common.MapStr{
				"@timestamp": common.Time(time.Now().UTC()),
				"type":       "pingbeat",
				"target": common.MapStr{
					"name": name,
					"addr": ping.Target,
					"tags": tags,
				},
				"loss":   true,
				"reason": ping.LossReason,
			}
			go bt.client.PublishEvent(event)
			logp.Debug("ProcessPing", "Processed ping error for %v (%v): %v", name, ping.Target, ping.LossReason)
		} else {
			event := common.MapStr{
				"@timestamp": common.Time(time.Now().UTC()),
				"type":       "pingbeat",
				"target": common.MapStr{
					"name": name,
					"addr": ping.Target,
					"tags": tags,
				},
				"rtt": milliSeconds(ping.RTT),
			}
			go bt.client.PublishEvent(event)
			logp.Debug("ProcessPing", "Processed ping %v for %v (%v): %v", ping.Seq, name, ping.Target, ping.RTT)
		}
	}
}

func parseICMPError(data []byte) (int, int, string) {
	IPheader, err := ipv4.ParseHeader(data[:len(data)-8])
	if err != nil {
		logp.Err("parseICMPError", "Failed to parse packet header:", err)
	}
	ICMPHdr := data[IPheader.Len:]
	var ID, Seq uint16
	err = binary.Read(bytes.NewReader(ICMPHdr[6:8]), binary.BigEndian, &Seq)
	if err != nil {
		logp.Err("parseICMPError", "Failed to parse packet header:", err)
	}
	err = binary.Read(bytes.NewReader(ICMPHdr[4:6]), binary.BigEndian, &ID)
	if err != nil {
		logp.Err("parseICMPError", "Failed to parse packet header:", err)
	}
	return int(ID), int(Seq), IPheader.Dst.String()
}

func createConn(n string, a string) (*icmp.PacketConn, error) {
	c, err := icmp.ListenPacket(n, a)
	if err != nil {
		return nil, err
	}
	return c, nil
}

func milliSeconds(d time.Duration) float64 {
	msec := d / time.Millisecond
	nsec := d % time.Millisecond
	return float64(msec) + float64(nsec)*1e-6
}
<|endoftext|>


<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 sdb

import (
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"testing"
)

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

////////////////////////////////////////////////////////////////////////
// PutAttributes
////////////////////////////////////////////////////////////////////////

type PutTest struct {
	domainTest

	item ItemName
	updates []PutUpdate
	preconditions []Precondition

	err error
}

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

func (t *PutTest) SetUp(i *TestInfo) {
	// Call common setup code.
	t.domainTest.SetUp(i)

	// Make the request legal by default.
	t.item = "foo"
	t.updates = []PutUpdate{PutUpdate{"bar", "baz", false}}
}

func (t *PutTest) callDomain() {
	t.err = t.domain.PutAttributes(t.item, t.updates, t.preconditions)
}

func (t *PutTest) EmptyItemName() {
	t.item = ""

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("item name")))
}

func (t *PutTest) InvalidItemName() {
	t.item = "taco\x80\x81\x82"

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("item name")))
	ExpectThat(t.err, Error(HasSubstr(string(t.item))))
}

func (t *PutTest) ZeroUpdates() {
	t.updates = []PutUpdate{}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("number")))
	ExpectThat(t.err, Error(HasSubstr("updates")))
	ExpectThat(t.err, Error(HasSubstr("0")))
}

func (t *PutTest) TooManyUpdates() {
	t.updates = make([]PutUpdate, 257)

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("number")))
	ExpectThat(t.err, Error(HasSubstr("updates")))
	ExpectThat(t.err, Error(HasSubstr("256")))
}

func (t *PutTest) OneAttributeNameEmpty() {
	t.updates = []PutUpdate{
		PutUpdate{Name: "foo"},
		PutUpdate{Name: "", Value: "taco"},
		PutUpdate{Name: "bar"},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *PutTest) OneAttributeNameInvalid() {
	t.updates = []PutUpdate{
		PutUpdate{Name: "foo"},
		PutUpdate{Name: "taco\x80\x81\x82"},
		PutUpdate{Name: "bar"},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
	ExpectThat(t.err, Error(HasSubstr(t.updates[1].Name)))
}

func (t *PutTest) OneAttributeValueInvalid() {
	t.updates = []PutUpdate{
		PutUpdate{Name: "foo"},
		PutUpdate{Name: "bar", Value: "taco\x80\x81\x82"},
		PutUpdate{Name: "baz"},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("value")))
	ExpectThat(t.err, Error(HasSubstr(t.updates[1].Value)))
}

func (t *PutTest) OnePreconditionNameEmpty() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "", Exists: new(bool)},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
}

func (t *PutTest) OnePreconditionNameInvalid() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "taco\x80\x81\x82", Exists: new(bool)},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
	ExpectThat(t.err, Error(HasSubstr(t.preconditions[1].Name)))
}

func (t *PutTest) OnePreconditionValueInvalid() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Value: new(string)},
		Precondition{Name: "bar", Value: new(string)},
		Precondition{Name: "baz", Value: new(string)},
	}

	*t.preconditions[0].Value = ""
	*t.preconditions[1].Value = "taco\x80\x81\x82"
	*t.preconditions[2].Value = "qux"

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("value")))
	ExpectThat(t.err, Error(HasSubstr(*t.preconditions[1].Value)))
}

func (t *PutTest) OnePreconditionMissingOperand() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "bar"},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("precondition")))
	ExpectThat(t.err, Error(HasSubstr("bar")))
}

func (t *PutTest) OnePreconditionHasTwoOperands() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "bar", Exists: new(bool), Value: new(string)},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("precondition")))
	ExpectThat(t.err, Error(HasSubstr("bar")))
}

func (t *PutTest) NoPreconditions() {
	ExpectEq("TODO", "")
}

func (t *PutTest) SomePreconditions() {
	ExpectEq("TODO", "")
}

func (t *PutTest) ConnReturnsError() {
	ExpectEq("TODO", "")
}

func (t *PutTest) ConnSaysOkay() {
	ExpectEq("TODO", "")
}

////////////////////////////////////////////////////////////////////////
// BatchPutAttributes
////////////////////////////////////////////////////////////////////////

type BatchPutTest struct {
	domainTest
}

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

func (t *BatchPutTest) DoesFoo() {
	ExpectEq("TODO", "")
}
<commit_msg>PutTest.NoPreconditions<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 sdb

import (
	"github.com/jacobsa/aws/exp/sdb/conn"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"sort"
	"testing"
)

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

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

func getSortedKeys(r conn.Request) []string {
	result := sort.StringSlice{}
	for key, _ := range r {
		result = append(result, key)
	}

	sort.Sort(result)
	return result
}

////////////////////////////////////////////////////////////////////////
// PutAttributes
////////////////////////////////////////////////////////////////////////

type PutTest struct {
	domainTest

	item ItemName
	updates []PutUpdate
	preconditions []Precondition

	err error
}

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

func (t *PutTest) SetUp(i *TestInfo) {
	// Call common setup code.
	t.domainTest.SetUp(i)

	// Make the request legal by default.
	t.item = "foo"
	t.updates = []PutUpdate{PutUpdate{"bar", "baz", false}}
}

func (t *PutTest) callDomain() {
	t.err = t.domain.PutAttributes(t.item, t.updates, t.preconditions)
}

func (t *PutTest) EmptyItemName() {
	t.item = ""

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("item name")))
}

func (t *PutTest) InvalidItemName() {
	t.item = "taco\x80\x81\x82"

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("item name")))
	ExpectThat(t.err, Error(HasSubstr(string(t.item))))
}

func (t *PutTest) ZeroUpdates() {
	t.updates = []PutUpdate{}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("number")))
	ExpectThat(t.err, Error(HasSubstr("updates")))
	ExpectThat(t.err, Error(HasSubstr("0")))
}

func (t *PutTest) TooManyUpdates() {
	t.updates = make([]PutUpdate, 257)

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("number")))
	ExpectThat(t.err, Error(HasSubstr("updates")))
	ExpectThat(t.err, Error(HasSubstr("256")))
}

func (t *PutTest) OneAttributeNameEmpty() {
	t.updates = []PutUpdate{
		PutUpdate{Name: "foo"},
		PutUpdate{Name: "", Value: "taco"},
		PutUpdate{Name: "bar"},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *PutTest) OneAttributeNameInvalid() {
	t.updates = []PutUpdate{
		PutUpdate{Name: "foo"},
		PutUpdate{Name: "taco\x80\x81\x82"},
		PutUpdate{Name: "bar"},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
	ExpectThat(t.err, Error(HasSubstr(t.updates[1].Name)))
}

func (t *PutTest) OneAttributeValueInvalid() {
	t.updates = []PutUpdate{
		PutUpdate{Name: "foo"},
		PutUpdate{Name: "bar", Value: "taco\x80\x81\x82"},
		PutUpdate{Name: "baz"},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("value")))
	ExpectThat(t.err, Error(HasSubstr(t.updates[1].Value)))
}

func (t *PutTest) OnePreconditionNameEmpty() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "", Exists: new(bool)},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
}

func (t *PutTest) OnePreconditionNameInvalid() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "taco\x80\x81\x82", Exists: new(bool)},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
	ExpectThat(t.err, Error(HasSubstr(t.preconditions[1].Name)))
}

func (t *PutTest) OnePreconditionValueInvalid() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Value: new(string)},
		Precondition{Name: "bar", Value: new(string)},
		Precondition{Name: "baz", Value: new(string)},
	}

	*t.preconditions[0].Value = ""
	*t.preconditions[1].Value = "taco\x80\x81\x82"
	*t.preconditions[2].Value = "qux"

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("value")))
	ExpectThat(t.err, Error(HasSubstr(*t.preconditions[1].Value)))
}

func (t *PutTest) OnePreconditionMissingOperand() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "bar"},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("precondition")))
	ExpectThat(t.err, Error(HasSubstr("bar")))
}

func (t *PutTest) OnePreconditionHasTwoOperands() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "bar", Exists: new(bool), Value: new(string)},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("precondition")))
	ExpectThat(t.err, Error(HasSubstr("bar")))
}

func (t *PutTest) NoPreconditions() {
	t.item = "some_item"
	t.updates = []PutUpdate{
		PutUpdate{Name: "foo"},
		PutUpdate{Name: "bar", Value: "taco", Replace: true},
		PutUpdate{Name: "baz", Value: "burrito"},
	}

	// Call
	t.callDomain()
	AssertNe(nil, t.c.req)

	AssertThat(
		getSortedKeys(t.c.req),
		ElementsAre(
			"Attribute.1.Name",
			"Attribute.1.Value",
			"Attribute.2.Name",
			"Attribute.2.Replace",
			"Attribute.2.Value",
			"Attribute.3.Name",
			"Attribute.3.Value",
			"DomainName",
			"ItemName",
		),
	)

	ExpectEq("foo", t.c.req["Attribute.1.Name"])
	ExpectEq("bar", t.c.req["Attribute.2.Name"])
	ExpectEq("baz", t.c.req["Attribute.3.Name"])

	ExpectEq("", t.c.req["Attribute.1.Value"])
	ExpectEq("taco", t.c.req["Attribute.2.Value"])
	ExpectEq("burrito", t.c.req["Attribute.3.Value"])

	ExpectEq("true", t.c.req["Attribute.2.Replace"])

	ExpectEq("some_item", t.c.req["ItemName"])
	ExpectEq(t.name, t.c.req["DomainName"])
}

func (t *PutTest) SomePreconditions() {
	ExpectEq("TODO", "")
}

func (t *PutTest) ConnReturnsError() {
	ExpectEq("TODO", "")
}

func (t *PutTest) ConnSaysOkay() {
	ExpectEq("TODO", "")
}

////////////////////////////////////////////////////////////////////////
// BatchPutAttributes
////////////////////////////////////////////////////////////////////////

type BatchPutTest struct {
	domainTest
}

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

func (t *BatchPutTest) DoesFoo() {
	ExpectEq("TODO", "")
}
<|endoftext|>


<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 sdb

import (
	"errors"
	"github.com/jacobsa/aws/exp/sdb/conn"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"sort"
	"testing"
)

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

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

func getSortedKeys(r conn.Request) []string {
	result := sort.StringSlice{}
	for key, _ := range r {
		result = append(result, key)
	}

	sort.Sort(result)
	return result
}

////////////////////////////////////////////////////////////////////////
// PutAttributes
////////////////////////////////////////////////////////////////////////

type PutTest struct {
	domainTest

	item          ItemName
	updates       []PutUpdate
	preconditions []Precondition

	err error
}

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

func (t *PutTest) SetUp(i *TestInfo) {
	// Call common setup code.
	t.domainTest.SetUp(i)

	// Make the request legal by default.
	t.item = "foo"
	t.updates = []PutUpdate{PutUpdate{"bar", "baz", false}}
}

func (t *PutTest) callDomain() {
	t.err = t.domain.PutAttributes(t.item, t.updates, t.preconditions)
}

func (t *PutTest) EmptyItemName() {
	t.item = ""

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("item name")))
}

func (t *PutTest) InvalidItemName() {
	t.item = "taco\x80\x81\x82"

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("item name")))
	ExpectThat(t.err, Error(HasSubstr("UTF-8")))
}

func (t *PutTest) ZeroUpdates() {
	t.updates = []PutUpdate{}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("number")))
	ExpectThat(t.err, Error(HasSubstr("updates")))
	ExpectThat(t.err, Error(HasSubstr("0")))
}

func (t *PutTest) TooManyUpdates() {
	t.updates = make([]PutUpdate, 257)

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("number")))
	ExpectThat(t.err, Error(HasSubstr("updates")))
	ExpectThat(t.err, Error(HasSubstr("257")))
}

func (t *PutTest) OneAttributeNameEmpty() {
	t.updates = []PutUpdate{
		PutUpdate{Name: "foo"},
		PutUpdate{Name: "", Value: "taco"},
		PutUpdate{Name: "bar"},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *PutTest) OneAttributeNameInvalid() {
	t.updates = []PutUpdate{
		PutUpdate{Name: "foo"},
		PutUpdate{Name: "taco\x80\x81\x82"},
		PutUpdate{Name: "bar"},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
	ExpectThat(t.err, Error(HasSubstr(t.updates[1].Name)))
}

func (t *PutTest) OneAttributeValueInvalid() {
	t.updates = []PutUpdate{
		PutUpdate{Name: "foo"},
		PutUpdate{Name: "bar", Value: "taco\x80\x81\x82"},
		PutUpdate{Name: "baz"},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("value")))
	ExpectThat(t.err, Error(HasSubstr(t.updates[1].Value)))
}

func (t *PutTest) OnePreconditionNameEmpty() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "", Exists: new(bool)},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
}

func (t *PutTest) OnePreconditionNameInvalid() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "taco\x80\x81\x82", Exists: new(bool)},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
	ExpectThat(t.err, Error(HasSubstr(t.preconditions[1].Name)))
}

func (t *PutTest) OnePreconditionValueInvalid() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Value: new(string)},
		Precondition{Name: "bar", Value: new(string)},
		Precondition{Name: "baz", Value: new(string)},
	}

	*t.preconditions[0].Value = ""
	*t.preconditions[1].Value = "taco\x80\x81\x82"
	*t.preconditions[2].Value = "qux"

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("value")))
}

func (t *PutTest) OnePreconditionMissingOperand() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "bar"},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("precondition")))
	ExpectThat(t.err, Error(HasSubstr("bar")))
}

func (t *PutTest) OnePreconditionHasTwoOperands() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "bar", Exists: new(bool), Value: new(string)},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("precondition")))
	ExpectThat(t.err, Error(HasSubstr("bar")))
}

func (t *PutTest) BasicParameters() {
	t.item = "some_item"
	t.updates = []PutUpdate{
		PutUpdate{Name: "foo"},
		PutUpdate{Name: "bar", Value: "taco", Replace: true},
		PutUpdate{Name: "baz", Value: "burrito"},
	}

	// Call
	t.callDomain()
	AssertNe(nil, t.c.req)

	AssertThat(
		getSortedKeys(t.c.req),
		ElementsAre(
			"Attribute.1.Name",
			"Attribute.1.Value",
			"Attribute.2.Name",
			"Attribute.2.Replace",
			"Attribute.2.Value",
			"Attribute.3.Name",
			"Attribute.3.Value",
			"DomainName",
			"ItemName",
		),
	)

	ExpectEq("foo", t.c.req["Attribute.1.Name"])
	ExpectEq("bar", t.c.req["Attribute.2.Name"])
	ExpectEq("baz", t.c.req["Attribute.3.Name"])

	ExpectEq("", t.c.req["Attribute.1.Value"])
	ExpectEq("taco", t.c.req["Attribute.2.Value"])
	ExpectEq("burrito", t.c.req["Attribute.3.Value"])

	ExpectEq("true", t.c.req["Attribute.2.Replace"])

	ExpectEq("some_item", t.c.req["ItemName"])
	ExpectEq(t.name, t.c.req["DomainName"])
}

func (t *PutTest) NoPreconditions() {
	// Call
	t.callDomain()
	AssertNe(nil, t.c.req)

	ExpectThat(getSortedKeys(t.c.req), Not(Contains(HasSubstr("Expected"))))
}

func (t *PutTest) SomePreconditions() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "bar", Value: new(string)},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	*t.preconditions[0].Exists = false
	*t.preconditions[1].Value = "taco"
	*t.preconditions[2].Exists = true

	// Call
	t.callDomain()
	AssertNe(nil, t.c.req)

	AssertThat(
		getSortedKeys(t.c.req),
		AllOf(
			Contains("Expected.1.Name"),
			Contains("Expected.2.Name"),
			Contains("Expected.3.Name"),
			Contains("Expected.1.Exists"),
			Contains("Expected.2.Value"),
			Contains("Expected.3.Exists"),
		),
	)

	ExpectEq("foo", t.c.req["Expected.1.Name"])
	ExpectEq("bar", t.c.req["Expected.2.Name"])
	ExpectEq("baz", t.c.req["Expected.3.Name"])

	ExpectEq("false", t.c.req["Expected.1.Exists"])
	ExpectEq("taco", t.c.req["Expected.2.Value"])
	ExpectEq("true", t.c.req["Expected.3.Exists"])
}

func (t *PutTest) ConnReturnsError() {
	// Conn
	t.c.err = errors.New("taco")

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("SendRequest")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *PutTest) ConnSaysOkay() {
	// Conn
	t.c.resp = []byte{}

	// Call
	t.callDomain()

	ExpectEq(nil, t.err)
}

////////////////////////////////////////////////////////////////////////
// BatchPutAttributes
////////////////////////////////////////////////////////////////////////

type BatchPutTest struct {
	domainTest

	updates       map[ItemName][]PutUpdate

	err error
}

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

func (t *BatchPutTest) SetUp(i *TestInfo) {
	// Call common setup code.
	t.domainTest.SetUp(i)

	// Make the request legal by default.
	t.updates = map[ItemName][]PutUpdate{
		"some_item": []PutUpdate{
			PutUpdate{Name: "foo"},
		},
	}
}

func (t *BatchPutTest) callDomain() {
	t.err = t.domain.BatchPutAttributes(t.updates)
}

func (t *BatchPutTest) NoItems() {
	t.updates = map[ItemName][]PutUpdate{
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("number")))
	ExpectThat(t.err, Error(HasSubstr("items")))
}

func (t *BatchPutTest) TooManyItems() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) OneItemNameEmpty() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) OneItemNameInvalid() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) ZeroUpdatesForOneItem() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) TooManyUpdatesForOneItem() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) OneAttributeNameEmpty() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) OneAttributeNameInvalid() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) OneAttributeValueInvalid() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) CallsConn() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) ConnReturnsError() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) ConnSaysOkay() {
	ExpectEq("TODO", "")
}
<commit_msg>BatchPutTest.TooManyItems<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 sdb

import (
	"errors"
	"fmt"
	"github.com/jacobsa/aws/exp/sdb/conn"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"sort"
	"testing"
)

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

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

func getSortedKeys(r conn.Request) []string {
	result := sort.StringSlice{}
	for key, _ := range r {
		result = append(result, key)
	}

	sort.Sort(result)
	return result
}

////////////////////////////////////////////////////////////////////////
// PutAttributes
////////////////////////////////////////////////////////////////////////

type PutTest struct {
	domainTest

	item          ItemName
	updates       []PutUpdate
	preconditions []Precondition

	err error
}

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

func (t *PutTest) SetUp(i *TestInfo) {
	// Call common setup code.
	t.domainTest.SetUp(i)

	// Make the request legal by default.
	t.item = "foo"
	t.updates = []PutUpdate{PutUpdate{"bar", "baz", false}}
}

func (t *PutTest) callDomain() {
	t.err = t.domain.PutAttributes(t.item, t.updates, t.preconditions)
}

func (t *PutTest) EmptyItemName() {
	t.item = ""

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("item name")))
}

func (t *PutTest) InvalidItemName() {
	t.item = "taco\x80\x81\x82"

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("item name")))
	ExpectThat(t.err, Error(HasSubstr("UTF-8")))
}

func (t *PutTest) ZeroUpdates() {
	t.updates = []PutUpdate{}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("number")))
	ExpectThat(t.err, Error(HasSubstr("updates")))
	ExpectThat(t.err, Error(HasSubstr("0")))
}

func (t *PutTest) TooManyUpdates() {
	t.updates = make([]PutUpdate, 257)

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("number")))
	ExpectThat(t.err, Error(HasSubstr("updates")))
	ExpectThat(t.err, Error(HasSubstr("257")))
}

func (t *PutTest) OneAttributeNameEmpty() {
	t.updates = []PutUpdate{
		PutUpdate{Name: "foo"},
		PutUpdate{Name: "", Value: "taco"},
		PutUpdate{Name: "bar"},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *PutTest) OneAttributeNameInvalid() {
	t.updates = []PutUpdate{
		PutUpdate{Name: "foo"},
		PutUpdate{Name: "taco\x80\x81\x82"},
		PutUpdate{Name: "bar"},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
	ExpectThat(t.err, Error(HasSubstr(t.updates[1].Name)))
}

func (t *PutTest) OneAttributeValueInvalid() {
	t.updates = []PutUpdate{
		PutUpdate{Name: "foo"},
		PutUpdate{Name: "bar", Value: "taco\x80\x81\x82"},
		PutUpdate{Name: "baz"},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("value")))
	ExpectThat(t.err, Error(HasSubstr(t.updates[1].Value)))
}

func (t *PutTest) OnePreconditionNameEmpty() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "", Exists: new(bool)},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
}

func (t *PutTest) OnePreconditionNameInvalid() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "taco\x80\x81\x82", Exists: new(bool)},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("name")))
	ExpectThat(t.err, Error(HasSubstr(t.preconditions[1].Name)))
}

func (t *PutTest) OnePreconditionValueInvalid() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Value: new(string)},
		Precondition{Name: "bar", Value: new(string)},
		Precondition{Name: "baz", Value: new(string)},
	}

	*t.preconditions[0].Value = ""
	*t.preconditions[1].Value = "taco\x80\x81\x82"
	*t.preconditions[2].Value = "qux"

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("attribute")))
	ExpectThat(t.err, Error(HasSubstr("value")))
}

func (t *PutTest) OnePreconditionMissingOperand() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "bar"},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("precondition")))
	ExpectThat(t.err, Error(HasSubstr("bar")))
}

func (t *PutTest) OnePreconditionHasTwoOperands() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "bar", Exists: new(bool), Value: new(string)},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("Invalid")))
	ExpectThat(t.err, Error(HasSubstr("precondition")))
	ExpectThat(t.err, Error(HasSubstr("bar")))
}

func (t *PutTest) BasicParameters() {
	t.item = "some_item"
	t.updates = []PutUpdate{
		PutUpdate{Name: "foo"},
		PutUpdate{Name: "bar", Value: "taco", Replace: true},
		PutUpdate{Name: "baz", Value: "burrito"},
	}

	// Call
	t.callDomain()
	AssertNe(nil, t.c.req)

	AssertThat(
		getSortedKeys(t.c.req),
		ElementsAre(
			"Attribute.1.Name",
			"Attribute.1.Value",
			"Attribute.2.Name",
			"Attribute.2.Replace",
			"Attribute.2.Value",
			"Attribute.3.Name",
			"Attribute.3.Value",
			"DomainName",
			"ItemName",
		),
	)

	ExpectEq("foo", t.c.req["Attribute.1.Name"])
	ExpectEq("bar", t.c.req["Attribute.2.Name"])
	ExpectEq("baz", t.c.req["Attribute.3.Name"])

	ExpectEq("", t.c.req["Attribute.1.Value"])
	ExpectEq("taco", t.c.req["Attribute.2.Value"])
	ExpectEq("burrito", t.c.req["Attribute.3.Value"])

	ExpectEq("true", t.c.req["Attribute.2.Replace"])

	ExpectEq("some_item", t.c.req["ItemName"])
	ExpectEq(t.name, t.c.req["DomainName"])
}

func (t *PutTest) NoPreconditions() {
	// Call
	t.callDomain()
	AssertNe(nil, t.c.req)

	ExpectThat(getSortedKeys(t.c.req), Not(Contains(HasSubstr("Expected"))))
}

func (t *PutTest) SomePreconditions() {
	t.preconditions = []Precondition{
		Precondition{Name: "foo", Exists: new(bool)},
		Precondition{Name: "bar", Value: new(string)},
		Precondition{Name: "baz", Exists: new(bool)},
	}

	*t.preconditions[0].Exists = false
	*t.preconditions[1].Value = "taco"
	*t.preconditions[2].Exists = true

	// Call
	t.callDomain()
	AssertNe(nil, t.c.req)

	AssertThat(
		getSortedKeys(t.c.req),
		AllOf(
			Contains("Expected.1.Name"),
			Contains("Expected.2.Name"),
			Contains("Expected.3.Name"),
			Contains("Expected.1.Exists"),
			Contains("Expected.2.Value"),
			Contains("Expected.3.Exists"),
		),
	)

	ExpectEq("foo", t.c.req["Expected.1.Name"])
	ExpectEq("bar", t.c.req["Expected.2.Name"])
	ExpectEq("baz", t.c.req["Expected.3.Name"])

	ExpectEq("false", t.c.req["Expected.1.Exists"])
	ExpectEq("taco", t.c.req["Expected.2.Value"])
	ExpectEq("true", t.c.req["Expected.3.Exists"])
}

func (t *PutTest) ConnReturnsError() {
	// Conn
	t.c.err = errors.New("taco")

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("SendRequest")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *PutTest) ConnSaysOkay() {
	// Conn
	t.c.resp = []byte{}

	// Call
	t.callDomain()

	ExpectEq(nil, t.err)
}

////////////////////////////////////////////////////////////////////////
// BatchPutAttributes
////////////////////////////////////////////////////////////////////////

type BatchPutTest struct {
	domainTest

	updates       map[ItemName][]PutUpdate

	err error
}

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

func (t *BatchPutTest) SetUp(i *TestInfo) {
	// Call common setup code.
	t.domainTest.SetUp(i)

	// Make the request legal by default.
	t.updates = map[ItemName][]PutUpdate{
		"some_item": []PutUpdate{
			PutUpdate{Name: "foo"},
		},
	}
}

func (t *BatchPutTest) callDomain() {
	t.err = t.domain.BatchPutAttributes(t.updates)
}

func (t *BatchPutTest) NoItems() {
	t.updates = map[ItemName][]PutUpdate{
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("number")))
	ExpectThat(t.err, Error(HasSubstr("items")))
	ExpectThat(t.err, Error(HasSubstr("0")))
}

func (t *BatchPutTest) TooManyItems() {
	t.updates = map[ItemName][]PutUpdate{}

	for i := 0; i < 26; i++ {
		t.updates[ItemName(fmt.Sprintf("%d", i))] = []PutUpdate{
			PutUpdate{Name: "foo"},
		}
	}

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("number")))
	ExpectThat(t.err, Error(HasSubstr("items")))
	ExpectThat(t.err, Error(HasSubstr("26")))
}

func (t *BatchPutTest) OneItemNameEmpty() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) OneItemNameInvalid() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) ZeroUpdatesForOneItem() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) TooManyUpdatesForOneItem() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) OneAttributeNameEmpty() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) OneAttributeNameInvalid() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) OneAttributeValueInvalid() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) CallsConn() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) ConnReturnsError() {
	ExpectEq("TODO", "")
}

func (t *BatchPutTest) ConnSaysOkay() {
	ExpectEq("TODO", "")
}
<|endoftext|>


<commit_before>package client

import (
	"encoding/binary"
	"fmt"
	capn "github.com/glycerine/go-capnproto"
	"goshawkdb.io/common"
	cmsgs "goshawkdb.io/common/capnp"
	"goshawkdb.io/server"
	msgs "goshawkdb.io/server/capnp"
	"goshawkdb.io/server/paxos"
	eng "goshawkdb.io/server/txnengine"
	"time"
)

type ClientTxnCompletionConsumer func(*cmsgs.ClientTxnOutcome, error)

type ClientTxnSubmitter struct {
	*SimpleTxnSubmitter
	versionCache versionCache
	txnLive      bool
	initialDelay time.Duration
}

func NewClientTxnSubmitter(rmId common.RMId, bootCount uint32, cm paxos.ConnectionManager) *ClientTxnSubmitter {
	return &ClientTxnSubmitter{
		SimpleTxnSubmitter: NewSimpleTxnSubmitter(rmId, bootCount, cm),
		versionCache:       NewVersionCache(),
		txnLive:            false,
		initialDelay:       time.Duration(0),
	}
}

func (cts *ClientTxnSubmitter) Status(sc *server.StatusConsumer) {
	sc.Emit(fmt.Sprintf("ClientTxnSubmitter: txnLive? %v", cts.txnLive))
	cts.SimpleTxnSubmitter.Status(sc.Fork())
	sc.Join()
}

func (cts *ClientTxnSubmitter) SubmitClientTransaction(ctxnCap *cmsgs.ClientTxn, continuation ClientTxnCompletionConsumer) {
	if cts.txnLive {
		continuation(nil, fmt.Errorf("Cannot submit client as a live txn already exists"))
		return
	}

	seg := capn.NewBuffer(nil)
	clientOutcome := cmsgs.NewClientTxnOutcome(seg)
	clientOutcome.SetId(ctxnCap.Id())

	curTxnId := common.MakeTxnId(ctxnCap.Id())

	delay := cts.initialDelay
	if delay < time.Millisecond {
		delay = time.Duration(0)
	}
	start := time.Now()

	var cont TxnCompletionConsumer
	cont = func(txn *eng.TxnReader, outcome *msgs.Outcome, err error) {
		if outcome == nil || err != nil { // node is shutting down or error
			cts.txnLive = false
			continuation(nil, err)
			return
		}
		txnId := txn.Id
		end := time.Now()
		elapsed := end.Sub(start)
		start = end
		switch outcome.Which() {
		case msgs.OUTCOME_COMMIT:
			cts.versionCache.UpdateFromCommit(txn, outcome)
			clientOutcome.SetFinalId(txnId[:])
			clientOutcome.SetCommit()
			cts.addCreatesToCache(txn)
			cts.txnLive = false
			cts.initialDelay = delay >> 1
			continuation(&clientOutcome, nil)
			return

		default:
			abort := outcome.Abort()
			resubmit := abort.Which() == msgs.OUTCOMEABORT_RESUBMIT
			if !resubmit {
				updates := abort.Rerun()
				validUpdates := cts.versionCache.UpdateFromAbort(&updates)
				server.Log("Updates:", updates.Len(), "; valid: ", len(validUpdates))
				resubmit = len(validUpdates) == 0
				if !resubmit {
					clientOutcome.SetFinalId(txnId[:])
					clientOutcome.SetAbort(cts.translateUpdates(seg, validUpdates))
					cts.txnLive = false
					cts.initialDelay = delay >> 1
					continuation(&clientOutcome, nil)
					return
				}
			}
			server.Log("Resubmitting", txnId, "; orig resubmit?", abort.Which() == msgs.OUTCOMEABORT_RESUBMIT)

			delay = delay + time.Duration(cts.rng.Intn(int(elapsed)))
			if delay > server.SubmissionMaxSubmitDelay {
				delay = server.SubmissionMaxSubmitDelay + time.Duration(cts.rng.Intn(int(server.SubmissionMaxSubmitDelay)))
			}
			//fmt.Printf("%v ", delay)

			curTxnIdNum := binary.BigEndian.Uint64(txnId[:8])
			curTxnIdNum += 1 + uint64(cts.rng.Intn(8))
			binary.BigEndian.PutUint64(curTxnId[:8], curTxnIdNum)
			ctxnCap.SetId(curTxnId[:])

			cts.SimpleTxnSubmitter.SubmitClientTransaction(ctxnCap, curTxnId, cont, delay, false)
		}
	}

	cts.txnLive = true
	// fmt.Printf("%v ", delay)
	cts.SimpleTxnSubmitter.SubmitClientTransaction(ctxnCap, curTxnId, cont, delay, false)
}

func (cts *ClientTxnSubmitter) addCreatesToCache(txn *eng.TxnReader) {
	actions := txn.Actions(true).Actions()
	for idx, l := 0, actions.Len(); idx < l; idx++ {
		action := actions.At(idx)
		if action.Which() == msgs.ACTION_CREATE {
			varUUId := common.MakeVarUUId(action.VarId())
			positions := common.Positions(action.Create().Positions())
			cts.hashCache.AddPosition(varUUId, &positions)
		}
	}
}

func (cts *ClientTxnSubmitter) translateUpdates(seg *capn.Segment, updates map[*msgs.Update][]*msgs.Action) cmsgs.ClientUpdate_List {
	clientUpdates := cmsgs.NewClientUpdateList(seg, len(updates))
	idx := 0
	for update, actions := range updates {
		clientUpdate := clientUpdates.At(idx)
		idx++
		clientUpdate.SetVersion(update.TxnId())
		clientActions := cmsgs.NewClientActionList(seg, len(actions))
		clientUpdate.SetActions(clientActions)

		for idy, action := range actions {
			clientAction := clientActions.At(idy)
			clientAction.SetVarId(action.VarId())
			switch action.Which() {
			case msgs.ACTION_MISSING:
				clientAction.SetDelete()
			case msgs.ACTION_WRITE:
				clientAction.SetWrite()
				write := action.Write()
				clientWrite := clientAction.Write()
				clientWrite.SetValue(write.Value())
				references := write.References()
				clientReferences := seg.NewDataList(references.Len())
				clientWrite.SetReferences(clientReferences)
				for idz, n := 0, references.Len(); idz < n; idz++ {
					ref := references.At(idz)
					clientReferences.Set(idz, ref.Id())
					positions := common.Positions(ref.Positions())
					cts.hashCache.AddPosition(common.MakeVarUUId(ref.Id()), &positions)
				}
			default:
				panic(fmt.Sprintf("Unexpected action type: %v", action.Which()))
			}
		}
	}
	return clientUpdates
}
<commit_msg>don't clobber capnp struct from client<commit_after>package client

import (
	"encoding/binary"
	"fmt"
	capn "github.com/glycerine/go-capnproto"
	"goshawkdb.io/common"
	cmsgs "goshawkdb.io/common/capnp"
	"goshawkdb.io/server"
	msgs "goshawkdb.io/server/capnp"
	"goshawkdb.io/server/paxos"
	eng "goshawkdb.io/server/txnengine"
	"time"
)

type ClientTxnCompletionConsumer func(*cmsgs.ClientTxnOutcome, error)

type ClientTxnSubmitter struct {
	*SimpleTxnSubmitter
	versionCache versionCache
	txnLive      bool
	initialDelay time.Duration
}

func NewClientTxnSubmitter(rmId common.RMId, bootCount uint32, cm paxos.ConnectionManager) *ClientTxnSubmitter {
	return &ClientTxnSubmitter{
		SimpleTxnSubmitter: NewSimpleTxnSubmitter(rmId, bootCount, cm),
		versionCache:       NewVersionCache(),
		txnLive:            false,
		initialDelay:       time.Duration(0),
	}
}

func (cts *ClientTxnSubmitter) Status(sc *server.StatusConsumer) {
	sc.Emit(fmt.Sprintf("ClientTxnSubmitter: txnLive? %v", cts.txnLive))
	cts.SimpleTxnSubmitter.Status(sc.Fork())
	sc.Join()
}

func (cts *ClientTxnSubmitter) SubmitClientTransaction(ctxnCap *cmsgs.ClientTxn, continuation ClientTxnCompletionConsumer) {
	if cts.txnLive {
		continuation(nil, fmt.Errorf("Cannot submit client as a live txn already exists"))
		return
	}

	seg := capn.NewBuffer(nil)
	clientOutcome := cmsgs.NewClientTxnOutcome(seg)
	clientOutcome.SetId(ctxnCap.Id())

	curTxnId := common.MakeTxnId(ctxnCap.Id())

	delay := cts.initialDelay
	if delay < time.Millisecond {
		delay = time.Duration(0)
	}
	start := time.Now()

	var cont TxnCompletionConsumer
	cont = func(txn *eng.TxnReader, outcome *msgs.Outcome, err error) {
		if outcome == nil || err != nil { // node is shutting down or error
			cts.txnLive = false
			continuation(nil, err)
			return
		}
		txnId := txn.Id
		end := time.Now()
		elapsed := end.Sub(start)
		start = end
		switch outcome.Which() {
		case msgs.OUTCOME_COMMIT:
			cts.versionCache.UpdateFromCommit(txn, outcome)
			clientOutcome.SetFinalId(txnId[:])
			clientOutcome.SetCommit()
			cts.addCreatesToCache(txn)
			cts.txnLive = false
			cts.initialDelay = delay >> 1
			continuation(&clientOutcome, nil)
			return

		default:
			abort := outcome.Abort()
			resubmit := abort.Which() == msgs.OUTCOMEABORT_RESUBMIT
			if !resubmit {
				updates := abort.Rerun()
				validUpdates := cts.versionCache.UpdateFromAbort(&updates)
				server.Log("Updates:", updates.Len(), "; valid: ", len(validUpdates))
				resubmit = len(validUpdates) == 0
				if !resubmit {
					clientOutcome.SetFinalId(txnId[:])
					clientOutcome.SetAbort(cts.translateUpdates(seg, validUpdates))
					cts.txnLive = false
					cts.initialDelay = delay >> 1
					continuation(&clientOutcome, nil)
					return
				}
			}
			server.Log("Resubmitting", txnId, "; orig resubmit?", abort.Which() == msgs.OUTCOMEABORT_RESUBMIT)

			delay = delay + time.Duration(cts.rng.Intn(int(elapsed)))
			if delay > server.SubmissionMaxSubmitDelay {
				delay = server.SubmissionMaxSubmitDelay + time.Duration(cts.rng.Intn(int(server.SubmissionMaxSubmitDelay)))
			}
			//fmt.Printf("%v ", delay)

			curTxnIdNum := binary.BigEndian.Uint64(txnId[:8])
			curTxnIdNum += 1 + uint64(cts.rng.Intn(8))
			binary.BigEndian.PutUint64(curTxnId[:8], curTxnIdNum)
			newSeg := capn.NewBuffer(nil)
			newCtxnCap := cmsgs.NewClientTxn(newSeg)
			newCtxnCap.SetId(curTxnId[:])
			newCtxnCap.SetRetry(ctxnCap.Retry())
			newCtxnCap.SetActions(ctxnCap.Actions())

			cts.SimpleTxnSubmitter.SubmitClientTransaction(&newCtxnCap, curTxnId, cont, delay, false)
		}
	}

	cts.txnLive = true
	// fmt.Printf("%v ", delay)
	cts.SimpleTxnSubmitter.SubmitClientTransaction(ctxnCap, curTxnId, cont, delay, false)
}

func (cts *ClientTxnSubmitter) addCreatesToCache(txn *eng.TxnReader) {
	actions := txn.Actions(true).Actions()
	for idx, l := 0, actions.Len(); idx < l; idx++ {
		action := actions.At(idx)
		if action.Which() == msgs.ACTION_CREATE {
			varUUId := common.MakeVarUUId(action.VarId())
			positions := common.Positions(action.Create().Positions())
			cts.hashCache.AddPosition(varUUId, &positions)
		}
	}
}

func (cts *ClientTxnSubmitter) translateUpdates(seg *capn.Segment, updates map[*msgs.Update][]*msgs.Action) cmsgs.ClientUpdate_List {
	clientUpdates := cmsgs.NewClientUpdateList(seg, len(updates))
	idx := 0
	for update, actions := range updates {
		clientUpdate := clientUpdates.At(idx)
		idx++
		clientUpdate.SetVersion(update.TxnId())
		clientActions := cmsgs.NewClientActionList(seg, len(actions))
		clientUpdate.SetActions(clientActions)

		for idy, action := range actions {
			clientAction := clientActions.At(idy)
			clientAction.SetVarId(action.VarId())
			switch action.Which() {
			case msgs.ACTION_MISSING:
				clientAction.SetDelete()
			case msgs.ACTION_WRITE:
				clientAction.SetWrite()
				write := action.Write()
				clientWrite := clientAction.Write()
				clientWrite.SetValue(write.Value())
				references := write.References()
				clientReferences := seg.NewDataList(references.Len())
				clientWrite.SetReferences(clientReferences)
				for idz, n := 0, references.Len(); idz < n; idz++ {
					ref := references.At(idz)
					clientReferences.Set(idz, ref.Id())
					positions := common.Positions(ref.Positions())
					cts.hashCache.AddPosition(common.MakeVarUUId(ref.Id()), &positions)
				}
			default:
				panic(fmt.Sprintf("Unexpected action type: %v", action.Which()))
			}
		}
	}
	return clientUpdates
}
<|endoftext|>


<commit_before>package command

import (
	"fmt"

	"github.com/google/yamlfmt"
	"github.com/google/yamlfmt/engine"
	"github.com/google/yamlfmt/formatters/basic"
	"github.com/mitchellh/mapstructure"
)

type Operation int

const (
	OperationFormat Operation = iota
	OperationLint
	OperationDry
)

type commandConfig struct {
	Include         []string               `mapstructure:"include"`
	Exclude         []string               `mapstructure:"exclude"`
	FormatterConfig map[string]interface{} `mapstructure:"formatter"`
}

func RunCommand(
	operation Operation,
	registry *yamlfmt.Registry,
	configData map[string]interface{},
) error {
	var config commandConfig
	err := mapstructure.Decode(configData, &config)
	if err != nil {
		return err
	}
	if len(config.Include) == 0 {
		config.Include = []string{"**/*.{yaml,yml}"}
	}

	var formatter yamlfmt.Formatter
	if config.FormatterConfig == nil {
		factory, err := registry.GetDefaultFactory()
		if err != nil {
			return err
		}
		formatter = factory.NewDefault()
		if err != nil {
			return err
		}
	} else {
		fType, ok := config.FormatterConfig["type"].(string)
		if !ok {
			fType = basic.BasicFormatterType
		}
		factory, err := registry.GetFactory(fType)
		if err != nil {
			return err
		}
		formatter, err = factory.NewWithConfig(config.FormatterConfig)
		if err != nil {
			return err
		}
	}

	engine := &engine.Engine{
		Include:   config.Include,
		Exclude:   config.Exclude,
		Formatter: formatter,
	}

	switch operation {
	case OperationFormat:
		err := engine.FormatAllFiles()
		if err != nil {
			return err
		}
	case OperationLint:
		err := engine.LintAllFiles()
		if err != nil {
			return err
		}
	case OperationDry:
		out, err := engine.DryRunAllFiles()
		if err != nil {
			return err
		}
		fmt.Println(out)
	}

	return nil
}
<commit_msg>delete unnecessary 'return err' (#1)<commit_after>package command

import (
	"fmt"

	"github.com/google/yamlfmt"
	"github.com/google/yamlfmt/engine"
	"github.com/google/yamlfmt/formatters/basic"
	"github.com/mitchellh/mapstructure"
)

type Operation int

const (
	OperationFormat Operation = iota
	OperationLint
	OperationDry
)

type commandConfig struct {
	Include         []string               `mapstructure:"include"`
	Exclude         []string               `mapstructure:"exclude"`
	FormatterConfig map[string]interface{} `mapstructure:"formatter"`
}

func RunCommand(
	operation Operation,
	registry *yamlfmt.Registry,
	configData map[string]interface{},
) error {
	var config commandConfig
	err := mapstructure.Decode(configData, &config)
	if err != nil {
		return err
	}
	if len(config.Include) == 0 {
		config.Include = []string{"**/*.{yaml,yml}"}
	}

	var formatter yamlfmt.Formatter
	if config.FormatterConfig == nil {
		factory, err := registry.GetDefaultFactory()
		if err != nil {
			return err
		}
		formatter = factory.NewDefault()
	} else {
		fType, ok := config.FormatterConfig["type"].(string)
		if !ok {
			fType = basic.BasicFormatterType
		}
		factory, err := registry.GetFactory(fType)
		if err != nil {
			return err
		}
		formatter, err = factory.NewWithConfig(config.FormatterConfig)
		if err != nil {
			return err
		}
	}

	engine := &engine.Engine{
		Include:   config.Include,
		Exclude:   config.Exclude,
		Formatter: formatter,
	}

	switch operation {
	case OperationFormat:
		err := engine.FormatAllFiles()
		if err != nil {
			return err
		}
	case OperationLint:
		err := engine.LintAllFiles()
		if err != nil {
			return err
		}
	case OperationDry:
		out, err := engine.DryRunAllFiles()
		if err != nil {
			return err
		}
		fmt.Println(out)
	}

	return nil
}
<|endoftext|>


<commit_before>package command

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

const (
	opSum   = "sum"
	opCount = "count"
	opMin   = "min"
	opMax   = "max"
	opMean  = "mean"

	predicateEq = "eq"
	predicateLt = "lt"
	predicateGt = "gt"
	predicateLe = "le"
	predicateGe = "ge"

	timeFmtSec    = "sec"
	timeFmtMsec   = "msec"
	timeFmtUsec   = "usec"
	timeFmtStruct = "timeval"

	outputJson = "json"
	outputCsv  = "csv"

	maxDurationStr = "24h"
)

var predicates = []string{predicateEq, predicateLt, predicateGt, predicateLe, predicateGe}
var ops = []string{opSum, opCount, opMin, opMax, opMean}
var timeFmts = []string{timeFmtSec, timeFmtMsec, timeFmtUsec, timeFmtStruct}
var outputs = []string{outputJson, outputCsv}

type exportData struct {
	projectId uint64
	namespace string
	fields    setFlags
	time      string
	wheres    setFlags

	groupBy  string
	operator string

	limitBy      string
	limitPeriods uint64
	limit        uint64

	timeFmt string

	output string

	rawQuery     string
	dumpRequest  bool
	dumpResponse bool
}

func (e *exportData) IsValid() bool {
	pidOk := e.projectId > 0
	limitOk := e.limit > 0

	opOk := len(e.operator) == 0 || isInList(e.operator, ops)

	groupOk := len(e.groupBy) == 0
	if !groupOk {
		if len(e.fields) != 1 {
			fmt.Println("There has to be exactly one field set when doing a groupBy.")
			return false
		}
		groupOk = len(e.operator) > 0 && isInList(e.operator, ops)
	}

	timeOk := isInList(e.timeFmt, timeFmts)
	outputOk := isInList(e.output, outputs)

	return pidOk && limitOk && opOk && groupOk && timeOk && outputOk
}

// NewExportCommand returns the base 'export' command.
func NewExportCommand(ctx *Context) *Command {
	e := new(exportData)
	pid := ctx.Profile.ActiveProject

	cmd := &Command{
		Name:    "query",
		ApiPath: "/v1/data",
		Usage:   "Get data for projects, devices, and fields.",
		Data:    e,
		Action:  getExport,
	}

	flags := cmd.NewFlagSet("iobeam query")

	flags.Uint64Var(&e.projectId, "projectId", pid, "Project ID (if omitted, defaults to active project)")
	flags.StringVar(&e.namespace, "namespace", "input", "Namespace to query.")
	flags.Var(&e.fields, "field", "Name of Field(s) to project results (flag can be used multiple times).")

	flags.StringVar(&e.time, "time", "", "Expects an interval from,to where from and to can be expressions like now()-2h or an absolute UNIX timestamp that defaults to milliseconds. The to-part is optional and defaults to now()")
	flags.Var(&e.wheres, "where", "A predicate statement on a field in the format f(field, value). Multiple where statements are allowed and form the logical conjunction (AND) (supported: "+strings.Join(predicates, ", ")+").")

	flags.StringVar(&e.groupBy, "groupBy", "", "requires the operator parameter to calculate the aggregate of a field over a specific time interval and by optional field. Ex. groupBy=time(2m),myField or groupBy=time(10s) where the time period can be ms, s, m, or h. Examples of valid values: '30s', '15m', '6h'. Requires a valid operator.")
	flags.StringVar(&e.operator, "operator", "", "Aggregation function to apply to datapoints: "+strings.Join(ops, ", "))

	flags.StringVar(&e.limitBy, "limitBy", "", "Max number of results per field (ex. location,10).")
	flags.Uint64Var(&e.limitPeriods, "limitPeriods", 0, "Limits the number of time periods to return in a group_by statement")
	flags.Uint64Var(&e.limit, "limit", 10, "Max number of results.")

	flags.StringVar(&e.timeFmt, "timeFmt", "msec", "Time unit to display timestamps: "+strings.Join(timeFmts, ", "))
	flags.StringVar(&e.output, "output", "json", "Output format of the results. Valid outputs: json, csv")

	flags.BoolVar(&e.dumpRequest, "dumpRequest", false, "Dump the request to std out.")
	flags.BoolVar(&e.dumpResponse, "dumpResponse", false, "Dump the response to std out.")

	return cmd
}

// getExport fetches the requested data from the iobeam Cloud based on
// the provided projectID, namespace, and fields names.
func getExport(c *Command, ctx *Context) error {
	e := c.Data.(*exportData)

	reqPath := c.ApiPath + "/" + e.namespace + "/"
	if len(e.fields) == 1 {
		for key := range e.fields {
			reqPath += key
			break
		}
	}

	req := ctx.Client.Get(reqPath).Expect(200).
		ProjectToken(ctx.Profile, e.projectId).
		DumpRequest(e.dumpRequest).
		DumpResponse(e.dumpResponse).
		ParamUint64("limit", e.limit).
		Param("timefmt", e.timeFmt).
		Param("output", e.output)

	if len(e.time) > 0 {
		req = req.Param("time", e.time)
	}

	if len(e.limitBy) > 0 {
		req = req.Param("limit_by", e.limitBy)
	}

	if len(e.operator) > 0 {
		req = req.Param("operator", e.operator)

		if len(e.groupBy) > 0 {
			req = req.Param("group_by", e.groupBy)
		}
	}

	if len(e.wheres) > 0 {
		for key := range e.wheres {
			req = req.Param("where", key)
		}
	}

	x := make(map[string]interface{})
	_, err := req.ResponseBody(&x).
		ResponseBodyHandler(func(body interface{}) error {

			if e.output == outputJson {
				if len(e.fields) > 0 && len(e.groupBy) == 0 {
					bodyMap := *(body.(*map[string]interface{}))
					results := bodyMap["result"].([]interface{})

					for _, v := range results {
						response := v.(map[string]interface{})
						fields := response["fields"].([]interface{})
						rows := response["values"].([]interface{})

						keepFields := make(map[string]bool)
						keepFieldsNames := make([]string, len(e.fields))
						keepIndexes := make(map[int]bool)

						for field, _ := range e.fields {
							keepFields[field] = true
						}

						j := 0
						for i, field := range fields {
							if keepFields[field.(string)] {
								keepIndexes[i] = true
								keepFieldsNames[j] = field.(string)
								j += 1
							} else {
								keepIndexes[i] = false
							}
						}

						filteredRows := make([]interface{}, len(rows))

						for rowIndex, row := range rows {
							filteredRow := make([]interface{}, len(keepFields))
							filteredRowIndex := 0
							for columnIndex, value := range row.([]interface{}) {
								if keepIndexes[columnIndex] {
									filteredRow[filteredRowIndex] = value
									filteredRowIndex += 1
								}
							}
							filteredRows[rowIndex] = filteredRow
						}

						response["fields"] = keepFieldsNames
						response["values"] = filteredRows
					}
				}

				output, err := json.MarshalIndent(body, "", "  ")
				fmt.Println(string(output))
				return err
			} else {
				var output string
				if len(e.fields) > 0 {

					keepFields := make(map[int]bool)

					for i, v := range strings.Split(body.(string), "\n") {
						if i == 0 {
							// This row is the column names

							for j, fieldName := range strings.Split(v, ",") {

								if j == 0 {
									// Timestamp field
									keepFields[j] = true
									output += fieldName
									continue
								}

								if _, ok := e.fields[strings.Split(fieldName, ".")[1]]; ok {
									// keep this column
									keepFields[j] = true
									output += "," + fieldName
								} else {
									keepFields[j] = false
								}
							}
							output += "\n"
							continue
						}

						for j, value := range strings.Split(v, ",") {
							if keepFields[j] == true {
								if output[len(output)-1:] != "\n" {
									output += ","
								}
								output += value

							}
						}
						output += "\n"
					}
				} else {
					output = body.(string)
				}
				fmt.Println(output)
			}
			return nil
		}).Execute()

	return err
}
<commit_msg>Fix limitPeriods in query<commit_after>package command

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

const (
	opSum   = "sum"
	opCount = "count"
	opMin   = "min"
	opMax   = "max"
	opMean  = "mean"

	predicateEq = "eq"
	predicateLt = "lt"
	predicateGt = "gt"
	predicateLe = "le"
	predicateGe = "ge"

	timeFmtSec    = "sec"
	timeFmtMsec   = "msec"
	timeFmtUsec   = "usec"
	timeFmtStruct = "timeval"

	outputJson = "json"
	outputCsv  = "csv"

	maxDurationStr = "24h"
)

var predicates = []string{predicateEq, predicateLt, predicateGt, predicateLe, predicateGe}
var ops = []string{opSum, opCount, opMin, opMax, opMean}
var timeFmts = []string{timeFmtSec, timeFmtMsec, timeFmtUsec, timeFmtStruct}
var outputs = []string{outputJson, outputCsv}

type exportData struct {
	projectId uint64
	namespace string
	fields    setFlags
	time      string
	wheres    setFlags

	groupBy  string
	operator string

	limitBy      string
	limitPeriods uint64
	limit        uint64

	timeFmt string

	output string

	rawQuery     string
	dumpRequest  bool
	dumpResponse bool
}

func (e *exportData) IsValid() bool {
	pidOk := e.projectId > 0
	limitOk := e.limit > 0

	opOk := len(e.operator) == 0 || isInList(e.operator, ops)

	groupOk := len(e.groupBy) == 0
	if !groupOk {
		if len(e.fields) != 1 {
			fmt.Println("There has to be exactly one field set when doing a groupBy.")
			return false
		}
		groupOk = len(e.operator) > 0 && isInList(e.operator, ops)
	}

	timeOk := isInList(e.timeFmt, timeFmts)
	outputOk := isInList(e.output, outputs)

	return pidOk && limitOk && opOk && groupOk && timeOk && outputOk
}

// NewExportCommand returns the base 'export' command.
func NewExportCommand(ctx *Context) *Command {
	e := new(exportData)
	pid := ctx.Profile.ActiveProject

	cmd := &Command{
		Name:    "query",
		ApiPath: "/v1/data",
		Usage:   "Get data for projects, devices, and fields.",
		Data:    e,
		Action:  getExport,
	}

	flags := cmd.NewFlagSet("iobeam query")

	flags.Uint64Var(&e.projectId, "projectId", pid, "Project ID (if omitted, defaults to active project)")
	flags.StringVar(&e.namespace, "namespace", "input", "Namespace to query.")
	flags.Var(&e.fields, "field", "Name of Field(s) to project results (flag can be used multiple times).")

	flags.StringVar(&e.time, "time", "", "Expects an interval from,to where from and to can be expressions like now()-2h or an absolute UNIX timestamp that defaults to milliseconds. The to-part is optional and defaults to now()")
	flags.Var(&e.wheres, "where", "A predicate statement on a field in the format f(field, value). Multiple where statements are allowed and form the logical conjunction (AND) (supported: "+strings.Join(predicates, ", ")+").")

	flags.StringVar(&e.groupBy, "groupBy", "", "requires the operator parameter to calculate the aggregate of a field over a specific time interval and by optional field. Ex. groupBy=time(2m),myField or groupBy=time(10s) where the time period can be ms, s, m, or h. Examples of valid values: '30s', '15m', '6h'. Requires a valid operator.")
	flags.StringVar(&e.operator, "operator", "", "Aggregation function to apply to datapoints: "+strings.Join(ops, ", "))

	flags.StringVar(&e.limitBy, "limitBy", "", "Max number of results per field (ex. location,10).")
	flags.Uint64Var(&e.limitPeriods, "limitPeriods", 0, "Limits the number of time periods to return in a group_by statement")
	flags.Uint64Var(&e.limit, "limit", 10, "Max number of results.")

	flags.StringVar(&e.timeFmt, "timeFmt", "msec", "Time unit to display timestamps: "+strings.Join(timeFmts, ", "))
	flags.StringVar(&e.output, "output", "json", "Output format of the results. Valid outputs: json, csv")

	flags.BoolVar(&e.dumpRequest, "dumpRequest", false, "Dump the request to std out.")
	flags.BoolVar(&e.dumpResponse, "dumpResponse", false, "Dump the response to std out.")

	return cmd
}

// getExport fetches the requested data from the iobeam Cloud based on
// the provided projectID, namespace, and fields names.
func getExport(c *Command, ctx *Context) error {
	e := c.Data.(*exportData)

	reqPath := c.ApiPath + "/" + e.namespace + "/"
	if len(e.fields) == 1 {
		for key := range e.fields {
			reqPath += key
			break
		}
	}

	req := ctx.Client.Get(reqPath).Expect(200).
		ProjectToken(ctx.Profile, e.projectId).
		DumpRequest(e.dumpRequest).
		DumpResponse(e.dumpResponse).
		ParamUint64("limit", e.limit).
		Param("timefmt", e.timeFmt).
		Param("output", e.output)

	if len(e.time) > 0 {
		req = req.Param("time", e.time)
	}

	if len(e.limitBy) > 0 {
		req = req.Param("limit_by", e.limitBy)
	}

	if e.limitPeriods > 0 {
		req = req.ParamUint64("limit_periods", e.limitPeriods)
	}

	if len(e.operator) > 0 {
		req = req.Param("operator", e.operator)

		if len(e.groupBy) > 0 {
			req = req.Param("group_by", e.groupBy)
		}
	}

	if len(e.wheres) > 0 {
		for key := range e.wheres {
			req = req.Param("where", key)
		}
	}

	x := make(map[string]interface{})
	_, err := req.ResponseBody(&x).
		ResponseBodyHandler(func(body interface{}) error {

			if e.output == outputJson {
				if len(e.fields) > 0 && len(e.groupBy) == 0 {
					bodyMap := *(body.(*map[string]interface{}))
					results := bodyMap["result"].([]interface{})

					for _, v := range results {
						response := v.(map[string]interface{})
						fields := response["fields"].([]interface{})
						rows := response["values"].([]interface{})

						keepFields := make(map[string]bool)
						keepFieldsNames := make([]string, len(e.fields))
						keepIndexes := make(map[int]bool)

						for field, _ := range e.fields {
							keepFields[field] = true
						}

						j := 0
						for i, field := range fields {
							if keepFields[field.(string)] {
								keepIndexes[i] = true
								keepFieldsNames[j] = field.(string)
								j += 1
							} else {
								keepIndexes[i] = false
							}
						}

						filteredRows := make([]interface{}, len(rows))

						for rowIndex, row := range rows {
							filteredRow := make([]interface{}, len(keepFields))
							filteredRowIndex := 0
							for columnIndex, value := range row.([]interface{}) {
								if keepIndexes[columnIndex] {
									filteredRow[filteredRowIndex] = value
									filteredRowIndex += 1
								}
							}
							filteredRows[rowIndex] = filteredRow
						}

						response["fields"] = keepFieldsNames
						response["values"] = filteredRows
					}
				}

				output, err := json.MarshalIndent(body, "", "  ")
				fmt.Println(string(output))
				return err
			} else {
				var output string
				if len(e.fields) > 0 {

					keepFields := make(map[int]bool)

					for i, v := range strings.Split(body.(string), "\n") {
						if i == 0 {
							// This row is the column names

							for j, fieldName := range strings.Split(v, ",") {

								if j == 0 {
									// Timestamp field
									keepFields[j] = true
									output += fieldName
									continue
								}

								if _, ok := e.fields[strings.Split(fieldName, ".")[1]]; ok {
									// keep this column
									keepFields[j] = true
									output += "," + fieldName
								} else {
									keepFields[j] = false
								}
							}
							output += "\n"
							continue
						}

						for j, value := range strings.Split(v, ",") {
							if keepFields[j] == true {
								if output[len(output)-1:] != "\n" {
									output += ","
								}
								output += value

							}
						}
						output += "\n"
					}
				} else {
					output = body.(string)
				}
				fmt.Println(output)
			}
			return nil
		}).Execute()

	return err
}
<|endoftext|>