<commit_before>package guestmetric

import (
	xenstoreclient "../xenstoreclient"
	"bufio"
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

type Collector struct {
	Client xenstoreclient.XenStoreClient
	Ballon bool
	Debug  bool
}

func (c *Collector) CollectOS() (GuestMetric, error) {
	current := make(GuestMetric, 0)
	f, err := os.OpenFile("/var/cache/xe-linux-distribution", os.O_RDONLY, 0666)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, "=") {
			parts := strings.SplitN(line, "=", 2)
			k := strings.TrimSpace(parts[0])
			v := strings.TrimSpace(strings.Trim(strings.TrimSpace(parts[1]), "\""))
			current[k] = v
		}
	}
	return prefixKeys("data/", current), nil
}

func (c *Collector) CollectMisc() (GuestMetric, error) {
	current := make(GuestMetric, 0)
	if c.Ballon {
		current["control/feature-balloon"] = "1"
	} else {
		current["control/feature-balloon"] = "0"
	}
	current["attr/PVAddons/Installed"] = "1"
	current["attr/PVAddons/MajorVersion"] = "@PRODUCT_MAJOR_VERSION@"
	current["attr/PVAddons/MinorVersion"] = "@PRODUCT_MINOR_VERSION@"
	current["attr/PVAddons/MicroVersion"] = "@PRODUCT_MICRO_VERSION@"
	current["attr/PVAddons/BuildVersion"] = "@NUMERIC_BUILD_NUMBER@"

	return current, nil
}

func (c *Collector) CollectMemory() (GuestMetric, error) {
	current := make(GuestMetric, 0)
	f, err := os.OpenFile("/proc/meminfo", os.O_RDONLY, 0666)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		parts := regexp.MustCompile(`\w+`).FindAllString(scanner.Text(), -1)
		switch parts[0] {
		case "MemTotal":
			current["meminfo_total"] = parts[1]
		case "MemFree":
			current["meminfo_free"] = parts[1]
		}
	}
	return prefixKeys("data/", current), nil
}

func enumNetworkAddresses(iface string) (GuestMetric, error) {
	const (
		IP_RE   string = `(\d{1,3}\.){3}\d{1,3}`
		IPV6_RE string = `[\da-f:]+[\da-f]`
	)

	var (
		IP_IPV4_ADDR_RE       = regexp.MustCompile(`inet\s*(` + IP_RE + `).*\se[a-zA-Z0-9]+[\s\n]`)
		IP_IPV6_ADDR_RE       = regexp.MustCompile(`inet6\s*(` + IPV6_RE + `)`)
		IFCONFIG_IPV4_ADDR_RE = regexp.MustCompile(`inet addr:\s*(` + IP_RE + `)`)
		IFCONFIG_IPV6_ADDR_RE = regexp.MustCompile(`inet6 addr:\s*(` + IPV6_RE + `)`)
	)

	d := make(GuestMetric, 0)

	var v4re, v6re *regexp.Regexp
	var out string
	var err error
	if out, err = runCmd("ip", "addr", "show", iface); err == nil {
		v4re = IP_IPV4_ADDR_RE
		v6re = IP_IPV6_ADDR_RE
	} else if out, err = runCmd("ifconfig", iface); err == nil {
		v4re = IFCONFIG_IPV4_ADDR_RE
		v6re = IFCONFIG_IPV6_ADDR_RE
	} else {
		return nil, fmt.Errorf("Cannot find ip/ifconfig command")
	}

	m := v4re.FindAllStringSubmatch(out, -1)
	if m != nil {
		for i, parts := range m {
			d[fmt.Sprintf("ipv4/%d", i)] = parts[1]
		}
	}
	m = v6re.FindAllStringSubmatch(out, -1)
	if m != nil {
		for i, parts := range m {
			d[fmt.Sprintf("ipv6/%d", i)] = parts[1]
		}
	}

	return d, nil
}

func getPlainVifId(path string) (string, error) {
	nodenamePath := fmt.Sprintf("%s/device/nodename", path)
	strLine, err := readSysfs(nodenamePath)
	if err != nil {
		return "", err
	}
	vifId := ""
	reNodename := regexp.MustCompile(`^device\/vif\/(\d+)$`)
	if matched := reNodename.FindStringSubmatch(strLine); matched != nil {
		vifId = matched[1]
	}
	if vifId == "" {
		return "", fmt.Errorf("Not found string like \"device/vif/[id]\" in file %s", nodenamePath)
	} else {
		return vifId, nil
	}
}

func (c *Collector) getSriovVifId(path string) (string, error) {
	sriovDevicePath := "xenserver/device/net-sriov-vf"
	macAddress, err := readSysfs(path + "/address")
	if err != nil {
		return "", err
	}
	subPaths, err := c.Client.List(sriovDevicePath)
	if err != nil {
		return "", err
	}
	for _, subPath := range subPaths {
		iterMac, err := c.Client.Read(fmt.Sprintf("%s/%s/mac", sriovDevicePath, subPath))
		if err != nil {
			continue
		}
		if iterMac == macAddress {
			return subPath, nil
		}
	}
	return "", fmt.Errorf("Cannot find a MAC address to map with %s", path)
}

// return vif_xenstore_prefix * vif_id * error where
// `vif_xenstore_prefix` could be either `attr/vif` for plain VIF or
// `xenserver/attr/net-sriov-vf` for SR-IOV VIF
func (c *Collector) getTargetXenstorePath(path string) (string, string, error) {
	plainVifPrefix := "attr/vif"
	sriovVifPrefix := "xenserver/attr/net-sriov-vf"
	// try to get `vif_id` from nodename interface, only a plain VIF have the nodename interface.
	vifId, err1 := getPlainVifId(path)
	if vifId != "" {
		return plainVifPrefix, vifId, nil
	}
	// not a plain VIF, it could possible be an SR-IOV VIF, try to get vif_id from MAC address mapping
	vifId, err2 := c.getSriovVifId(path)
	if vifId != "" {
		return sriovVifPrefix, vifId, nil
	}
	return "", "", fmt.Errorf("Failed to get VIF ID, errors: %s | %s", err1.Error(), err2.Error())
}

func (c *Collector) CollectNetworkAddr() (GuestMetric, error) {
	current := make(GuestMetric, 0)

	var paths []string
	vifNamePrefixList := [...]string{"eth", "eno", "ens", "emp", "enx"}
	for _, prefix := range vifNamePrefixList {
		prefixPaths, err := filepath.Glob(fmt.Sprintf("/sys/class/net/%s*", prefix))
		if err != nil {
			return nil, err
		}
		paths = append(paths, prefixPaths...)
	}
	for _, path := range paths {
		// a path is going to be like "/sys/class/net/eth0"
		prefix, vifId, err := c.getTargetXenstorePath(path)
		if err != nil {
			continue
		}
		iface := filepath.Base(path)
		if addrs, err := enumNetworkAddresses(iface); err == nil {
			for tag, addr := range addrs {
				current[fmt.Sprintf("%s/%s/%s", prefix, vifId, tag)] = addr
			}
		}
	}
	return current, nil
}

func readSysfs(filename string) (string, error) {
	f, err := os.OpenFile(filename, os.O_RDONLY, 0666)
	if err != nil {
		return "", err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	scanner.Scan()
	return scanner.Text(), nil
}

func (c *Collector) CollectDisk() (GuestMetric, error) {
	pi := make(GuestMetric, 0)

	disks := make([]string, 0)
	paths, err := filepath.Glob("/sys/block/*/device")
	if err != nil {
		return nil, err
	}
	for _, path := range paths {
		disk := filepath.Base(strings.TrimSuffix(filepath.Dir(path), "/"))
		disks = append(disks, disk)
	}

	var sortedDisks sort.StringSlice = disks
	sortedDisks.Sort()

	part_idx := 0
	for _, disk := range sortedDisks[:] {
		paths, err = filepath.Glob(fmt.Sprintf("/dev/%s?*", disk))
		if err != nil {
			return nil, err
		}
		for _, path := range paths {
			p := filepath.Base(path)
			line, err := readSysfs(fmt.Sprintf("/sys/block/%s/%s/size", disk, p))
			if err != nil {
				return nil, err
			}
			size, err := strconv.ParseInt(line, 10, 64)
			if err != nil {
				return nil, err
			}
			blocksize := 512
			if bs, err := readSysfs(fmt.Sprintf("/sys/block/%s/queue/physical_block_size", p)); err == nil {
				if bs1, err := strconv.Atoi(bs); err == nil {
					blocksize = bs1
				}
			}
			real_dev := ""
			if c.Client != nil {
				nodename, err := readSysfs(fmt.Sprintf("/sys/block/%s/device/nodename", disk))
				if err == nil {
					backend, err := c.Client.Read(fmt.Sprintf("%s/backend", nodename))
					if err != nil {
						return nil, err
					}
					real_dev, err = c.Client.Read(fmt.Sprintf("%s/dev", backend))
					if err != nil {
						return nil, err
					}
				}
			}
			name := path
			blkid, err := runCmd("blkid", "-s", "UUID", path)
			if err != nil {
				// ignore blkid errors
				blkid = ""
			}
			if strings.Contains(blkid, "=") {
				parts := strings.SplitN(strings.TrimSpace(blkid), "=", 2)
				name = fmt.Sprintf("%s(%s)", name, strings.Trim(parts[1], "\""))
			}
			i := map[string]string{
				"extents/0": real_dev,
				"name":      name,
				"size":      strconv.FormatInt(size*int64(blocksize), 10),
			}
			output, err := runCmd("pvs", "--noheadings", "--units", "b", path)
			if err == nil && output != "" {
				parts := regexp.MustCompile(`\s+`).Split(output, -1)[1:]
				i["free"] = strings.TrimSpace(parts[5])[:len(parts[5])-1]
				i["filesystem"] = strings.TrimSpace(parts[2])
				i["mount_points/0"] = "[LVM]"
			} else {
				output, err = runCmd("mount")
				if err == nil {
					m := regexp.MustCompile(`(?m)^(\S+) on (\S+) type (\S+)`).FindAllStringSubmatch(output, -1)
					if m != nil {
						for _, parts := range m {
							if parts[1] == path {
								i["mount_points/0"] = parts[2]
								i["filesystem"] = parts[3]
								break
							}
						}
					}
				}
				output, err = runCmd("df", path)
				if err == nil {
					scanner := bufio.NewScanner(bytes.NewReader([]byte(output)))
					scanner.Scan()
					scanner.Scan()
					parts := regexp.MustCompile(`\s+`).Split(scanner.Text(), -1)
					free, err := strconv.ParseInt(parts[3], 10, 64)
					if err == nil {
						i["free"] = strconv.FormatInt(free*1024, 10)
					}
				}
			}
			for k, v := range i {
				pi[fmt.Sprintf("data/volumes/%d/%s", part_idx, k)] = v
			}
			part_idx += 1
		}
	}
	return pi, nil
}
<commit_msg>CA-352995: Remove the unnecessary flag and permission when open file<commit_after>package guestmetric

import (
	xenstoreclient "../xenstoreclient"
	"bufio"
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

type Collector struct {
	Client xenstoreclient.XenStoreClient
	Ballon bool
	Debug  bool
}

func (c *Collector) CollectOS() (GuestMetric, error) {
	current := make(GuestMetric, 0)
	f, err := os.OpenFile("/var/cache/xe-linux-distribution", os.O_RDONLY, 0666)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, "=") {
			parts := strings.SplitN(line, "=", 2)
			k := strings.TrimSpace(parts[0])
			v := strings.TrimSpace(strings.Trim(strings.TrimSpace(parts[1]), "\""))
			current[k] = v
		}
	}
	return prefixKeys("data/", current), nil
}

func (c *Collector) CollectMisc() (GuestMetric, error) {
	current := make(GuestMetric, 0)
	if c.Ballon {
		current["control/feature-balloon"] = "1"
	} else {
		current["control/feature-balloon"] = "0"
	}
	current["attr/PVAddons/Installed"] = "1"
	current["attr/PVAddons/MajorVersion"] = "@PRODUCT_MAJOR_VERSION@"
	current["attr/PVAddons/MinorVersion"] = "@PRODUCT_MINOR_VERSION@"
	current["attr/PVAddons/MicroVersion"] = "@PRODUCT_MICRO_VERSION@"
	current["attr/PVAddons/BuildVersion"] = "@NUMERIC_BUILD_NUMBER@"

	return current, nil
}

func (c *Collector) CollectMemory() (GuestMetric, error) {
	current := make(GuestMetric, 0)
	f, err := os.Open("/proc/meminfo")
	if err != nil {
		return nil, err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		parts := regexp.MustCompile(`\w+`).FindAllString(scanner.Text(), -1)
		switch parts[0] {
		case "MemTotal":
			current["meminfo_total"] = parts[1]
		case "MemFree":
			current["meminfo_free"] = parts[1]
		}
	}
	return prefixKeys("data/", current), nil
}

func enumNetworkAddresses(iface string) (GuestMetric, error) {
	const (
		IP_RE   string = `(\d{1,3}\.){3}\d{1,3}`
		IPV6_RE string = `[\da-f:]+[\da-f]`
	)

	var (
		IP_IPV4_ADDR_RE       = regexp.MustCompile(`inet\s*(` + IP_RE + `).*\se[a-zA-Z0-9]+[\s\n]`)
		IP_IPV6_ADDR_RE       = regexp.MustCompile(`inet6\s*(` + IPV6_RE + `)`)
		IFCONFIG_IPV4_ADDR_RE = regexp.MustCompile(`inet addr:\s*(` + IP_RE + `)`)
		IFCONFIG_IPV6_ADDR_RE = regexp.MustCompile(`inet6 addr:\s*(` + IPV6_RE + `)`)
	)

	d := make(GuestMetric, 0)

	var v4re, v6re *regexp.Regexp
	var out string
	var err error
	if out, err = runCmd("ip", "addr", "show", iface); err == nil {
		v4re = IP_IPV4_ADDR_RE
		v6re = IP_IPV6_ADDR_RE
	} else if out, err = runCmd("ifconfig", iface); err == nil {
		v4re = IFCONFIG_IPV4_ADDR_RE
		v6re = IFCONFIG_IPV6_ADDR_RE
	} else {
		return nil, fmt.Errorf("Cannot find ip/ifconfig command")
	}

	m := v4re.FindAllStringSubmatch(out, -1)
	if m != nil {
		for i, parts := range m {
			d[fmt.Sprintf("ipv4/%d", i)] = parts[1]
		}
	}
	m = v6re.FindAllStringSubmatch(out, -1)
	if m != nil {
		for i, parts := range m {
			d[fmt.Sprintf("ipv6/%d", i)] = parts[1]
		}
	}

	return d, nil
}

func getPlainVifId(path string) (string, error) {
	nodenamePath := fmt.Sprintf("%s/device/nodename", path)
	strLine, err := readSysfs(nodenamePath)
	if err != nil {
		return "", err
	}
	vifId := ""
	reNodename := regexp.MustCompile(`^device\/vif\/(\d+)$`)
	if matched := reNodename.FindStringSubmatch(strLine); matched != nil {
		vifId = matched[1]
	}
	if vifId == "" {
		return "", fmt.Errorf("Not found string like \"device/vif/[id]\" in file %s", nodenamePath)
	} else {
		return vifId, nil
	}
}

func (c *Collector) getSriovVifId(path string) (string, error) {
	sriovDevicePath := "xenserver/device/net-sriov-vf"
	macAddress, err := readSysfs(path + "/address")
	if err != nil {
		return "", err
	}
	subPaths, err := c.Client.List(sriovDevicePath)
	if err != nil {
		return "", err
	}
	for _, subPath := range subPaths {
		iterMac, err := c.Client.Read(fmt.Sprintf("%s/%s/mac", sriovDevicePath, subPath))
		if err != nil {
			continue
		}
		if iterMac == macAddress {
			return subPath, nil
		}
	}
	return "", fmt.Errorf("Cannot find a MAC address to map with %s", path)
}

// return vif_xenstore_prefix * vif_id * error where
// `vif_xenstore_prefix` could be either `attr/vif` for plain VIF or
// `xenserver/attr/net-sriov-vf` for SR-IOV VIF
func (c *Collector) getTargetXenstorePath(path string) (string, string, error) {
	plainVifPrefix := "attr/vif"
	sriovVifPrefix := "xenserver/attr/net-sriov-vf"
	// try to get `vif_id` from nodename interface, only a plain VIF have the nodename interface.
	vifId, err1 := getPlainVifId(path)
	if vifId != "" {
		return plainVifPrefix, vifId, nil
	}
	// not a plain VIF, it could possible be an SR-IOV VIF, try to get vif_id from MAC address mapping
	vifId, err2 := c.getSriovVifId(path)
	if vifId != "" {
		return sriovVifPrefix, vifId, nil
	}
	return "", "", fmt.Errorf("Failed to get VIF ID, errors: %s | %s", err1.Error(), err2.Error())
}

func (c *Collector) CollectNetworkAddr() (GuestMetric, error) {
	current := make(GuestMetric, 0)

	var paths []string
	vifNamePrefixList := [...]string{"eth", "eno", "ens", "emp", "enx"}
	for _, prefix := range vifNamePrefixList {
		prefixPaths, err := filepath.Glob(fmt.Sprintf("/sys/class/net/%s*", prefix))
		if err != nil {
			return nil, err
		}
		paths = append(paths, prefixPaths...)
	}
	for _, path := range paths {
		// a path is going to be like "/sys/class/net/eth0"
		prefix, vifId, err := c.getTargetXenstorePath(path)
		if err != nil {
			continue
		}
		iface := filepath.Base(path)
		if addrs, err := enumNetworkAddresses(iface); err == nil {
			for tag, addr := range addrs {
				current[fmt.Sprintf("%s/%s/%s", prefix, vifId, tag)] = addr
			}
		}
	}
	return current, nil
}

func readSysfs(filename string) (string, error) {
	f, err := os.OpenFile(filename, os.O_RDONLY, 0666)
	if err != nil {
		return "", err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	scanner.Scan()
	return scanner.Text(), nil
}

func (c *Collector) CollectDisk() (GuestMetric, error) {
	pi := make(GuestMetric, 0)

	disks := make([]string, 0)
	paths, err := filepath.Glob("/sys/block/*/device")
	if err != nil {
		return nil, err
	}
	for _, path := range paths {
		disk := filepath.Base(strings.TrimSuffix(filepath.Dir(path), "/"))
		disks = append(disks, disk)
	}

	var sortedDisks sort.StringSlice = disks
	sortedDisks.Sort()

	part_idx := 0
	for _, disk := range sortedDisks[:] {
		paths, err = filepath.Glob(fmt.Sprintf("/dev/%s?*", disk))
		if err != nil {
			return nil, err
		}
		for _, path := range paths {
			p := filepath.Base(path)
			line, err := readSysfs(fmt.Sprintf("/sys/block/%s/%s/size", disk, p))
			if err != nil {
				return nil, err
			}
			size, err := strconv.ParseInt(line, 10, 64)
			if err != nil {
				return nil, err
			}
			blocksize := 512
			if bs, err := readSysfs(fmt.Sprintf("/sys/block/%s/queue/physical_block_size", p)); err == nil {
				if bs1, err := strconv.Atoi(bs); err == nil {
					blocksize = bs1
				}
			}
			real_dev := ""
			if c.Client != nil {
				nodename, err := readSysfs(fmt.Sprintf("/sys/block/%s/device/nodename", disk))
				if err == nil {
					backend, err := c.Client.Read(fmt.Sprintf("%s/backend", nodename))
					if err != nil {
						return nil, err
					}
					real_dev, err = c.Client.Read(fmt.Sprintf("%s/dev", backend))
					if err != nil {
						return nil, err
					}
				}
			}
			name := path
			blkid, err := runCmd("blkid", "-s", "UUID", path)
			if err != nil {
				// ignore blkid errors
				blkid = ""
			}
			if strings.Contains(blkid, "=") {
				parts := strings.SplitN(strings.TrimSpace(blkid), "=", 2)
				name = fmt.Sprintf("%s(%s)", name, strings.Trim(parts[1], "\""))
			}
			i := map[string]string{
				"extents/0": real_dev,
				"name":      name,
				"size":      strconv.FormatInt(size*int64(blocksize), 10),
			}
			output, err := runCmd("pvs", "--noheadings", "--units", "b", path)
			if err == nil && output != "" {
				parts := regexp.MustCompile(`\s+`).Split(output, -1)[1:]
				i["free"] = strings.TrimSpace(parts[5])[:len(parts[5])-1]
				i["filesystem"] = strings.TrimSpace(parts[2])
				i["mount_points/0"] = "[LVM]"
			} else {
				output, err = runCmd("mount")
				if err == nil {
					m := regexp.MustCompile(`(?m)^(\S+) on (\S+) type (\S+)`).FindAllStringSubmatch(output, -1)
					if m != nil {
						for _, parts := range m {
							if parts[1] == path {
								i["mount_points/0"] = parts[2]
								i["filesystem"] = parts[3]
								break
							}
						}
					}
				}
				output, err = runCmd("df", path)
				if err == nil {
					scanner := bufio.NewScanner(bytes.NewReader([]byte(output)))
					scanner.Scan()
					scanner.Scan()
					parts := regexp.MustCompile(`\s+`).Split(scanner.Text(), -1)
					free, err := strconv.ParseInt(parts[3], 10, 64)
					if err == nil {
						i["free"] = strconv.FormatInt(free*1024, 10)
					}
				}
			}
			for k, v := range i {
				pi[fmt.Sprintf("data/volumes/%d/%s", part_idx, k)] = v
			}
			part_idx += 1
		}
	}
	return pi, nil
}
<|endoftext|>


<commit_before>/*
	HEROTS
	HERald Of The Swarm

	Package for fast TLS server creation.
*/
package herots

import (
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
)

type Server struct {
	options *Options
	certs   struct {
		Cert tls.Certificate
		pool struct {
			IsSet bool
			Pool  *x509.CertPool
		}
	}
	listener       net.Listener
	logDestination io.Writer
}

// predefined errors messages
const (
	LoadKeyPairError      = "herots: load key pair error"
	LoadClientCaCertError = "herots srv: load client CA cert error"
	StartServerError      = "herots srv: start tls server error"
	NoKeyPairLoad         = "herots: no load key pair (use LoadKeyPair func)"
	AcceptConnError       = "herots srv: connection accept error"
)

/*
	A Options structure is used to configure a TLS server.
*/
type Options struct {
	// Server host.
	// By default server use "127.0.0.1".
	Host string

	// Server port.
	// By default server use "9000".
	Port int

	// LogLevel provides the opportunity to choose the level of
	// information messages.
	// Each level includes the messages from the previous level.
	// 0 - no messages
	// 1 - notice
	// 2 - info
	// 3 - error
	//
	// Default: '0'.
	LogLevel int

	// LogDestination provides the opportunity to choose the own
	// destination for log messages (errors, info, etc).
	//
	// Default: 'os.Stdout'.
	LogDestination io.Writer

	// See http://golang.org/pkg/crypto/tls/#ClientAuthType
	// By default server use tls.RequireAnyClientCert
	TLSAuthType tls.ClientAuthType
}

/*
	Return Server struct with predefined options.
*/
func NewServer(o *Options) *Server {
	s := &Server{}

	// check mandatory options
	if o.LogDestination == nil {
		o.LogDestination = os.Stdout
	}

	if o.Port == 0 {
		o.Port = 9000
	}

	if o.TLSAuthType == 0 {
		o.TLSAuthType = tls.RequireAnyClientCert
	}

	s.options = o

	return s
}

// func for print messages
func (h *Server) log(m string, lvl int) {
	if h.options.LogLevel == 0 {
		return
	}

	if h.options.LogLevel <= lvl {
		fmt.Fprintf(h.logDestination, "herots srv: %s\n", m)
	}
}

/*
	Func for load certificate and private key pair.

	Public/private key pair require as PEM encoded data.
*/
func (h *Server) LoadKeyPair(cert, key []byte) error {
	// create cert pool
	h.certs.pool.Pool = x509.NewCertPool()

	// load keypair
	c, err := tls.X509KeyPair(cert, key)
	if err != nil {
		return fmt.Errorf("%s: %v\n", LoadKeyPairError, err)
	}
	h.certs.Cert = c

	// add cert to pool
	pemData, _ := pem.Decode(cert)
	ca, err := x509.ParseCertificate(pemData.Bytes)
	if err != nil {
		return fmt.Errorf("%s: %v\n", LoadKeyPairError, err)
	}
	h.certs.pool.Pool.AddCert(ca)
	h.certs.pool.IsSet = true

	h.log("load key pair ok", 2)

	return nil
}

/*
	Add client CA certificate to x509.CertPool (tls.Config.ClientCAs).

	By default server add cert from server public/private key pair (LoadKeyPair)
	to cert pool.
*/
func (h *Server) AddClientCACert(cert []byte) error {
	pemData, _ := pem.Decode(cert)
	ca, err := x509.ParseCertificate(pemData.Bytes)
	if err != nil {
		return fmt.Errorf("%s: %v\n", LoadClientCaCertError, err)
	}
	h.certs.pool.Pool.AddCert(ca)

	h.log("load client CA cert ok", 2)

	return nil
}

/*
	Accept and return connection to server.
*/
func (h *Server) Accept() (net.Conn, error) {
	conn, err := h.listener.Accept()
	if err != nil {
		h.log("accept conn error: "+err.Error(), 3)
		return conn, fmt.Errorf("%s: %v\n", AcceptConnError, err)
	}
	h.log("accepted conn from "+conn.RemoteAddr().String(), 2)
	return conn, nil
}

/*
	Start server.
*/
func (h *Server) Start() error {
	// load keypair check
	if len(h.certs.Cert.Certificate) == 0 {
		return fmt.Errorf("%s\n", NoKeyPairLoad)
	}

	config := tls.Config{
		ClientAuth:   h.options.TLSAuthType,
		Certificates: []tls.Certificate{h.certs.Cert},
		ClientCAs:    h.certs.pool.Pool,
		Rand:         rand.Reader,
	}

	service := h.options.Host + ":" + strconv.Itoa(h.options.Port)

	listener, err := tls.Listen("tcp", service, &config)
	if err != nil {
		return fmt.Errorf("%s: %v\n", StartServerError, err)
	}
	h.listener = listener

	h.log("listening on "+service, 2)

	return nil
}
<commit_msg>change prefix name for better readability<commit_after>/*
	HEROTS
	HERald Of The Swarm

	Package for fast TLS server creation.
*/
package herots

import (
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
)

type Server struct {
	options *Options
	certs   struct {
		Cert tls.Certificate
		pool struct {
			IsSet bool
			Pool  *x509.CertPool
		}
	}
	listener       net.Listener
	logDestination io.Writer
}

// predefined errors messages
const (
	LoadKeyPairError      = "herots: load key pair error"
	LoadClientCaCertError = "herots srv: load client CA cert error"
	StartServerError      = "herots srv: start tls server error"
	NoKeyPairLoad         = "herots: no load key pair (use LoadKeyPair func)"
	AcceptConnError       = "herots srv: connection accept error"
)

/*
	A Options structure is used to configure a TLS server.
*/
type Options struct {
	// Server host.
	// By default server use "127.0.0.1".
	Host string

	// Server port.
	// By default server use "9000".
	Port int

	// LogLevel provides the opportunity to choose the level of
	// information messages.
	// Each level includes the messages from the previous level.
	// 0 - no messages
	// 1 - notice
	// 2 - info
	// 3 - error
	//
	// Default: '0'.
	LogLevel int

	// LogDestination provides the opportunity to choose the own
	// destination for log messages (errors, info, etc).
	//
	// Default: 'os.Stdout'.
	LogDestination io.Writer

	// See http://golang.org/pkg/crypto/tls/#ClientAuthType
	// By default server use tls.RequireAnyClientCert
	TLSAuthType tls.ClientAuthType
}

/*
	Return Server struct with predefined options.
*/
func NewServer(o *Options) *Server {
	s := &Server{}

	// check mandatory options
	if o.LogDestination == nil {
		o.LogDestination = os.Stdout
	}

	if o.Port == 0 {
		o.Port = 9000
	}

	if o.TLSAuthType == 0 {
		o.TLSAuthType = tls.RequireAnyClientCert
	}

	s.options = o

	return s
}

// func for print messages
func (s *Server) log(m string, lvl int) {
	if s.options.LogLevel == 0 {
		return
	}

	if s.options.LogLevel <= lvl {
		fmt.Fprintf(s.logDestination, "herots srv: %s\n", m)
	}
}

/*
	Func for load certificate and private key pair.

	Public/private key pair require as PEM encoded data.
*/
func (s *Server) LoadKeyPair(cert, key []byte) error {
	// create cert pool
	s.certs.pool.Pool = x509.NewCertPool()

	// load keypair
	c, err := tls.X509KeyPair(cert, key)
	if err != nil {
		return fmt.Errorf("%s: %v\n", LoadKeyPairError, err)
	}
	s.certs.Cert = c

	// add cert to pool
	pemData, _ := pem.Decode(cert)
	ca, err := x509.ParseCertificate(pemData.Bytes)
	if err != nil {
		return fmt.Errorf("%s: %v\n", LoadKeyPairError, err)
	}
	s.certs.pool.Pool.AddCert(ca)
	s.certs.pool.IsSet = true

	s.log("load key pair ok", 2)

	return nil
}

/*
	Add client CA certificate to x509.CertPool (tls.Config.ClientCAs).

	By default server add cert from server public/private key pair (LoadKeyPair)
	to cert pool.
*/
func (s *Server) AddClientCACert(cert []byte) error {
	pemData, _ := pem.Decode(cert)
	ca, err := x509.ParseCertificate(pemData.Bytes)
	if err != nil {
		return fmt.Errorf("%s: %v\n", LoadClientCaCertError, err)
	}
	s.certs.pool.Pool.AddCert(ca)

	s.log("load client CA cert ok", 2)

	return nil
}

/*
	Accept and return connection to server.
*/
func (s *Server) Accept() (net.Conn, error) {
	conn, err := s.listener.Accept()
	if err != nil {
		s.log("accept conn error: "+err.Error(), 3)
		return conn, fmt.Errorf("%s: %v\n", AcceptConnError, err)
	}
	s.log("accepted conn from "+conn.RemoteAddr().String(), 2)
	return conn, nil
}

/*
	Start server.
*/
func (s *Server) Start() error {
	// load keypair check
	if len(s.certs.Cert.Certificate) == 0 {
		return fmt.Errorf("%s\n", NoKeyPairLoad)
	}

	config := tls.Config{
		ClientAuth:   s.options.TLSAuthType,
		Certificates: []tls.Certificate{s.certs.Cert},
		ClientCAs:    s.certs.pool.Pool,
		Rand:         rand.Reader,
	}

	service := s.options.Host + ":" + strconv.Itoa(s.options.Port)

	listener, err := tls.Listen("tcp", service, &config)
	if err != nil {
		return fmt.Errorf("%s: %v\n", StartServerError, err)
	}
	s.listener = listener

	s.log("listening on "+service, 2)

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package executor

import (
	"github.com/cznic/mathutil"
	"github.com/pingcap/tidb/types"
	"github.com/pingcap/tidb/util/chunk"
	"golang.org/x/net/context"
)

// ExplainExec represents an explain executor.
type ExplainExec struct {
	baseExecutor

	rows   [][]string
	cursor int
}

// Next implements Execution Next interface.
func (e *ExplainExec) Next(ctx context.Context) (Row, error) {
	if e.cursor >= len(e.rows) {
		return nil, nil
	}
	resultRow := make([]types.Datum, 0, len(e.rows[0]))
	for i := range e.rows[e.cursor] {
		resultRow = append(resultRow, types.NewStringDatum(e.rows[e.cursor][i]))
	}
	e.cursor++
	return resultRow, nil
}

// Close implements the Executor Close interface.
func (e *ExplainExec) Close() error {
	e.rows = nil
	return nil
}

// NextChunk implements the Executor NextChunk interface.
func (e *ExplainExec) NextChunk(ctx context.Context, chk *chunk.Chunk) error {
	chk.Reset()
	if e.cursor >= len(e.rows) {
		return nil
	}

	numCurRows := mathutil.Min(e.maxChunkSize, len(e.rows)-e.cursor)
	for i := e.cursor; i < e.cursor+numCurRows; i++ {
		for j := range e.rows[i] {
			chk.AppendString(j, e.rows[i][j])
		}
	}
	e.cursor += numCurRows
	return nil
}
<commit_msg>executor: remove Next function for ExplainExec (#6002)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package executor

import (
	"github.com/cznic/mathutil"
	"github.com/pingcap/tidb/util/chunk"
	"golang.org/x/net/context"
)

// ExplainExec represents an explain executor.
type ExplainExec struct {
	baseExecutor

	rows   [][]string
	cursor int
}

// Close implements the Executor Close interface.
func (e *ExplainExec) Close() error {
	e.rows = nil
	return nil
}

// NextChunk implements the Executor NextChunk interface.
func (e *ExplainExec) NextChunk(ctx context.Context, chk *chunk.Chunk) error {
	chk.Reset()
	if e.cursor >= len(e.rows) {
		return nil
	}

	numCurRows := mathutil.Min(e.maxChunkSize, len(e.rows)-e.cursor)
	for i := e.cursor; i < e.cursor+numCurRows; i++ {
		for j := range e.rows[i] {
			chk.AppendString(j, e.rows[i][j])
		}
	}
	e.cursor += numCurRows
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2017, TCN Inc.
// All rights reserved.

// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:

//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of TCN Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.

// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package templates

const SpannerUnaryMethodTemplate = `{{define "spanner_unary_method"}}// spanner unary select {{.GetName}}
func (s* {{.GetServiceName}}Impl) {{.GetName}} (ctx context.Context, req *{{.GetInputType}}) (*{{.GetOutputType}}, error) {
{{if .Spanner.IsSelect}}{{template "spanner_unary_select" .}}{{end}}
{{if .Spanner.IsUpdate}}{{template "spanner_unary_update" .}}{{end}}
{{if .Spanner.IsInsert}}{{template "spanner_unary_insert" .}}{{end}}
{{if .Spanner.IsDelete}}{{template "spanner_unary_delete" .}}{{end}}
{{end}}`

const SpannerUnarySelectTemplate = `{{define "spanner_unary_select"}}
	var err error
	var (
{{range $field, $type := .GetFieldsWithLocalTypesFor .GetOutputTypeStruct}}
		{{$field}} {{$type}}{{end}}
	)

	{{template "declare_spanner_arg_map" .}}

	//stmt := spanner.Statement{SQL: "{ {.Spanner.Query} }", Params: params}
	stmt := spanner.Statement{SQL: "{{.Spanner.Query}}", Params: params}
	tx := s.SpannerDB.Single()
	defer tx.Close()
	iter := tx.Query(ctx, stmt)
	defer iter.Stop()
	row, err := iter.Next()
	if err == iterator.Done {
		return nil, grpc.Errorf(codes.NotFound, "no rows found")
	} else if err != nil {
		return nil, grpc.Errorf(codes.Unknown, err.Error())
	}

	// scan our values out of the row
	{{range $index, $t := .GetTypeDescArrayForStruct .GetOutputTypeStruct}}
	{{if $t.IsMapped}}
	gcv := new(spanner.GenericColumnValue)
	err = row.ColumnByName("{{$t.ProtoName}}", gcv)
	if err != nil {
		return nil, grpc.Errorf(codes.Unknown, err.Error())
	}
	err = {{$t.Name}}.SpannerScan(gcv)
	if err != nil {
		return nil, grpc.Errorf(codes.Unknown, err.Error())
	}
	{{else}}
	err = row.ColumnByName("{{$t.ProtoName}}", &{{$t.Name}})
	if err != nil {
		return nil, grpc.Errorf(codes.Unknown, err.Error())
	}
	{{end}}{{end}}

	_, err = iter.Next()
	if err != iterator.Done {
		fmt.Println("Unary select that returns more than one row..")
	}
	//res := &{ {.GetOutputType} }{
	res := &{{.GetOutputType}}{
	{{range $field, $type := .GetTypeDescForFieldsInStruct .GetOutputTypeStruct}}
	{{$field}}: {{template "addr" $type}}{{template "base" $type}}{{template "mapping" $type}},{{end}}
	}
	return res, nil
}
{{end}}`

const SpannerUnaryInsertTemplate = `{{define "spanner_unary_insert"}}
	var err error
	{{template "declare_spanner_arg_slice" .}}

	muts := make([]*spanner.Mutation, 1)
	muts[0] = spanner.Insert("{{.Spanner.TableName}}", {{.Spanner.InsertColsAsString}}, params)
	_, err = s.SpannerDB.Apply(ctx, muts)
	if err != nil {
		if strings.Contains(err.Error(), "already exists") {
			return nil, grpc.Errorf(codes.AlreadyExists, err.Error())
		} else {
			return nil, grpc.Errorf(codes.Unknown, err.Error())
		}
	}
	res := &{{.GetOutputType}}{}

	return res, nil
}
{{end}}`

const SpannerUnaryUpdateTemplate = `{{define "spanner_unary_update"}}
	var err error
	{{template "declare_spanner_arg_map" .}}

	muts := make([]*spanner.Mutation, 1)
	muts[0] = spanner.UpdateMap("{{.Spanner.TableName}}", params)
	_, err = s.SpannerDB.Apply(ctx, muts)
	if err != nil {
		if strings.Contains(err.Error(), "already exists") {
			return nil, grpc.Errorf(codes.AlreadyExists, err.Error())
		} else {
			return nil, grpc.Errorf(codes.Unknown, err.Error())
		}
	}
	res := &{{.GetOutputType}}{}

	return res, nil
}
{{end}}`

const SpannerUnaryDeleteTemplate = `{{define "spanner_unary_delete"}}
	var err error
{{template "declare_spanner_delete_key" .}}

	muts := make([]*spanner.Mutation, 1)
	muts[0] = spanner.DeleteKeyRange({{.Spanner.TableName}}, key)
	_, err = s.SpannerDB.Apply(ctx, muts)
	if err != nil {
		if strings.Contains(err.Error(), "does not exist") {
			return nil, grpc.Errorf(codes.NotFound, err.Error())
		}
	}
	res := &{{.GetOutputType}}{}

	return res, nil
}
{{end}}`

const SpannerClientStreamingMethodTemplate = `{{define "spanner_client_streaming_method"}}// spanner client streaming {{.GetName}}
func (s *{{.GetServiceName}}Impl) {{.GetName}}(stream {{.GetServiceName}}_{{.GetName}}Server) error {
	var totalAffected int64
	var err error
	muts := make([]*spanner.Mutation, 0)
	for {
		req, err := stream.Recv()
		if err == io.EOF {
			break
		} else if err  != nil {
			return grpc.Errorf(codes.Unknown, err.Error())
		}
		totalAffected += 1

		{{if .Spanner.IsInsert}}{{template "spanner_client_streaming_insert" .}}{{end}}
		{{if .Spanner.IsUpdate}}{{template "spanner_client_streaming_update" .}}{{end}}
		{{if .Spanner.IsDelete}}{{template "spanner_client_streaming_delete" .}}{{end}}
		//In the future, we might do apply if muts gets really big,  but for now,
		// we only do one apply on the database with all the records stored in muts
	}
	_, err = s.SpannerDB.Apply(context.Background(), muts)
	if err != nil {
		if strings.Contains(err.Error(), "already exists") {
			return grpc.Errorf(codes.AlreadyExists, err.Error())
		} else {
			return grpc.Errorf(codes.Unknown, err.Error())
		}
	}
	stream.SendAndClose(&{{.GetOutputType}}{Count: totalAffected})
	return nil
}
{{end}}`

const SpannerClientStreamingUpdateTemplate = `{{define "spanner_client_streaming_update"}}//spanner client streaming update
{{template "declare_spanner_arg_map" .}}

muts = append(muts, spanner.UpdateMap("{{.Spanner.TableName}}", params))
{{end}}`

const SpannerClientStreamingInsertTemplate = `{{define "spanner_client_streaming_insert"}}//spanner client streaming insert
{{template "declare_spanner_arg_slice" .}}

	muts = append(muts, spanner.Insert("{{.Spanner.TableName}}", {{.Spanner.InsertColsAsString}}, params))
{{end}}`

const SpannerClientStreamingDeleteTemplate = `{{define "spanner_client_streaming_delete"}}//spanner client streaming delete
{{template "declare_spanner_delete_key" .}}

	muts = append(muts, spanner.DeleteKeyRange({{.Spanner.TableName}}, key))
{{end}}`

const SpannerServerStreamingMethodTemplate = `{{define "spanner_server_streaming_method"}}// spanner server streaming {{.GetName}}
func (s *{{.GetServiceName}}Impl) {{.GetName}}(req *{{.GetInputType}}, stream {{.GetServiceName}}_{{.GetName}}Server) error {
	var (
	{{range $field, $type := .GetFieldsWithLocalTypesFor .GetOutputTypeStruct}}
		{{$field}} {{$type}}{{end}}
	)
	{{if ne (len .Spanner.QueryArgs) 0}}
	var err error
	{{end}}

	{{template "declare_spanner_arg_map" .}}

	stmt := spanner.Statement{SQL: "{{.Spanner.Query}}", Params: params}
	tx := s.SpannerDB.Single()
	defer tx.Close()
	iter := tx.Query(context.Background(), stmt)
	defer iter.Stop()
	for {
		row, err := iter.Next()
		if err == iterator.Done {
			break
		} else if err != nil {
			return grpc.Errorf(codes.Unknown, err.Error())
		}

		// scan our values out of the row
		{{range $index, $t := .GetTypeDescArrayForStruct .GetOutputTypeStruct}}
		{{if $t.IsMapped}}
		gcv := new(spanner.GenericColumnValue)
		err = row.ColumnByName("{{$t.ProtoName}}", gcv)
		if err != nil {
			return grpc.Errorf(codes.Unknown, err.Error())
		}
		err = {{$t.Name}}.SpannerScan(gcv)
		if err != nil {
			return grpc.Errorf(codes.Unknown, err.Error())
		}
		{{else}}
		err = row.ColumnByName("{{$t.ProtoName}}", &{{$t.Name}})
		if err != nil {
			return grpc.Errorf(codes.Unknown, err.Error())
		}
		{{end}}{{end}}
		res := &{{.GetOutputType}}{
		{{range $field, $type := .GetTypeDescForFieldsInStruct .GetOutputTypeStruct}}
		{{$field}}: {{template "addr" $type}}{{template "base" $type}}{{template "mapping" $type}},{{end}}
		}
		stream.Send(res)
	}
	return  nil
}
{{end}}`

const SpannerBidiStreamingMethodTemplate = `{{define "spanner_bidi_streaming_method"}}// spanner bidi streaming {{.GetName}}
unimplemented
{{end}}`

const SpannerHelperTemplates = `
{{define "type_desc_to_def_map"}}
{{if .IsMapped}}
	conv, err = {{.GoName}}{}.ToSpanner(req.{{.Name}}).SpannerValue()
{{else}}
	conv = req.{{.Name}}
{{end}}{{end}}

{{define "type_desc_to_def_slice"}}
{{if .IsMapped}}
	conv, err = {{.GoName}}{}.ToSpanner(req.{{.Name}}).SpannerValue()
{{else}}
	conv = req.{{.Name}}
{{end}}{{end}}

{{define "return_err_on_method"}}
	if err != nil {
{{if .IsUnary}}
		return nil, grpc.Errorf(codes.Unknown, err.Error())
{{else}}
		return grpc.Errorf(codes.Unknown, err.Error())
{{end}}
	}
{{end}}

{{define "declare_spanner_arg_map"}}
{{$method := .}}
	params := make(map[string]interface{})
{{if gt (len .Spanner.OptionArguments) 0}}
	var conv interface{}
{{end}}
{{range $key, $val := .Spanner.QueryArgs}}
{{if $val.IsFieldValue}}
	{{template "type_desc_to_def_map" $val.Field}}
	{{template "return_err_on_method" $method}}
	params["{{$val.Name}}"] = conv
{{else}}
{{if $val.IsValue}}
	conv = {{$val.Value}}
	params["{{$val.Name}}"] = conv
{{end}}
{{end}}{{end}}
{{end}}

{{define "declare_spanner_arg_slice"}}
{{$method := .}}
	params := make([]interface{}, 0)
{{if gt (len .Spanner.OptionArguments) 0}}
	var conv interface{}
{{end}}

{{range $index, $val := .Spanner.QueryArgs}}
{{if $val.IsFieldValue}}
	{{template "type_desc_to_def_slice" $val.Field}}
	{{template "return_err_on_method" $method}}
	params = append(params, conv)
{{else}}
{{if $val.IsValue}}
	params = append(params, {{$val.Value}})
{{end}}
{{end}}{{end}}
{{end}}

{{define "declare_spanner_delete_key"}}
{{$method := .}}
	start := make([]interface{}, 0)
	end := make([]interface{}, 0)
{{if gt (len .Spanner.OptionArguments) 0}}
	var conv interface{}
{{end}}
{{range $index, $arg := .Spanner.KeyRangeDesc.Start}}
{{if $arg.IsFieldValue}}
{{template "type_desc_to_def_slice" $arg.Field}}
{{template "return_err_on_method" $method}}
	start = append(start, conv)
{{else}}
	start = append(start, {{$arg.Value}})
{{end}}{{end}}
{{range $index, $arg := .Spanner.KeyRangeDesc.End}}
{{if $arg.IsFieldValue}}
{{template "type_desc_to_def_slice" $arg.Field}}
{{template "return_err_on_method" $method}}
	end = append(end, conv)
{{else}}
{{if $arg.IsValue}}
	end = append(end, {{$arg.Value}})
{{end}}
{{end}}{{end}}
	key := spanner.KeyRange{
		Start: start,
		End: end,
		Kind: spanner.{{.Spanner.KeyRangeDesc.Kind}},
	}
{{end}}

`
<commit_msg>spanner templates use before/after hooks<commit_after>// Copyright 2017, TCN Inc.
// All rights reserved.

// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:

//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of TCN Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.

// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package templates

const SpannerUnaryMethodTemplate = `{{define "spanner_unary_method"}}// spanner unary select {{.GetName}}
func (s* {{.GetServiceName}}Impl) {{.GetName}} (ctx context.Context, req *{{.GetInputType}}) (*{{.GetOutputType}}, error) {
{{if .Spanner.IsSelect}}{{template "spanner_unary_select" .}}{{end}}
{{if .Spanner.IsUpdate}}{{template "spanner_unary_update" .}}{{end}}
{{if .Spanner.IsInsert}}{{template "spanner_unary_insert" .}}{{end}}
{{if .Spanner.IsDelete}}{{template "spanner_unary_delete" .}}{{end}}
{{end}}`

const SpannerUnarySelectTemplate = `{{define "spanner_unary_select"}}
	var err error
	var (
{{range $field, $type := .GetFieldsWithLocalTypesFor .GetOutputTypeStruct}}
		{{$field}} {{$type}}{{end}}
	)

	{{template "before_hook" .}}
	{{template "declare_spanner_arg_map" .}}

	//stmt := spanner.Statement{SQL: "{ {.Spanner.Query} }", Params: params}
	stmt := spanner.Statement{SQL: "{{.Spanner.Query}}", Params: params}
	tx := s.SpannerDB.Single()
	defer tx.Close()
	iter := tx.Query(ctx, stmt)
	defer iter.Stop()
	row, err := iter.Next()
	if err == iterator.Done {
		return nil, grpc.Errorf(codes.NotFound, "no rows found")
	} else if err != nil {
		return nil, grpc.Errorf(codes.Unknown, err.Error())
	}

	// scan our values out of the row
	{{range $index, $t := .GetTypeDescArrayForStruct .GetOutputTypeStruct}}
	{{if $t.IsMapped}}
	gcv := new(spanner.GenericColumnValue)
	err = row.ColumnByName("{{$t.ProtoName}}", gcv)
	if err != nil {
		return nil, grpc.Errorf(codes.Unknown, err.Error())
	}
	err = {{$t.Name}}.SpannerScan(gcv)
	if err != nil {
		return nil, grpc.Errorf(codes.Unknown, err.Error())
	}
	{{else}}
	err = row.ColumnByName("{{$t.ProtoName}}", &{{$t.Name}})
	if err != nil {
		return nil, grpc.Errorf(codes.Unknown, err.Error())
	}
	{{end}}{{end}}

	_, err = iter.Next()
	if err != iterator.Done {
		fmt.Println("Unary select that returns more than one row..")
	}
	res := {{.GetOutputType}}{
	{{range $field, $type := .GetTypeDescForFieldsInStruct .GetOutputTypeStruct}}
	{{$field}}: {{template "addr" $type}}{{template "base" $type}}{{template "mapping" $type}},{{end}}
	}

	{{template "after_hook" .}}

	return &res, nil
}
{{end}}`

const SpannerUnaryInsertTemplate = `{{define "spanner_unary_insert"}}
	var err error
	{{template "before_hook" .}}
	{{template "declare_spanner_arg_slice" .}}

	muts := make([]*spanner.Mutation, 1)
	muts[0] = spanner.Insert("{{.Spanner.TableName}}", {{.Spanner.InsertColsAsString}}, params)
	_, err = s.SpannerDB.Apply(ctx, muts)
	if err != nil {
		if strings.Contains(err.Error(), "already exists") {
			return nil, grpc.Errorf(codes.AlreadyExists, err.Error())
		} else {
			return nil, grpc.Errorf(codes.Unknown, err.Error())
		}
	}
	res := {{.GetOutputType}}{}

	{{template "after_hook" .}}

	return &res, nil
}
{{end}}`

const SpannerUnaryUpdateTemplate = `{{define "spanner_unary_update"}}
	var err error

	{{template "before_hook" .}}
	{{template "declare_spanner_arg_map" .}}

	muts := make([]*spanner.Mutation, 1)
	muts[0] = spanner.UpdateMap("{{.Spanner.TableName}}", params)
	_, err = s.SpannerDB.Apply(ctx, muts)
	if err != nil {
		if strings.Contains(err.Error(), "already exists") {
			return nil, grpc.Errorf(codes.AlreadyExists, err.Error())
		} else {
			return nil, grpc.Errorf(codes.Unknown, err.Error())
		}
	}
	res := {{.GetOutputType}}{}

	{{template "after_hook" .}}

	return &res, nil
}
{{end}}`

const SpannerUnaryDeleteTemplate = `{{define "spanner_unary_delete"}}
	var err error

	{{template "before_hook" .}}
	{{template "declare_spanner_delete_key" .}}

	muts := make([]*spanner.Mutation, 1)
	muts[0] = spanner.DeleteKeyRange({{.Spanner.TableName}}, key)
	_, err = s.SpannerDB.Apply(ctx, muts)
	if err != nil {
		if strings.Contains(err.Error(), "does not exist") {
			return nil, grpc.Errorf(codes.NotFound, err.Error())
		}
	}
	res := {{.GetOutputType}}{}

	{{template "after_hook" .}}

	return &res, nil
}
{{end}}`

const SpannerClientStreamingMethodTemplate = `{{define "spanner_client_streaming_method"}}// spanner client streaming {{.GetName}}
func (s *{{.GetServiceName}}Impl) {{.GetName}}(stream {{.GetServiceName}}_{{.GetName}}Server) error {
	var err error
	res := {{.GetOutputType}}{}
	muts := make([]*spanner.Mutation, 0)
	for {
		req, err := stream.Recv()
		if err == io.EOF {
			break
		} else if err  != nil {
			return grpc.Errorf(codes.Unknown, err.Error())
		}

		{{template "before_hook" .}}
		{{if .Spanner.IsInsert}}{{template "spanner_client_streaming_insert" .}}{{end}}
		{{if .Spanner.IsUpdate}}{{template "spanner_client_streaming_update" .}}{{end}}
		{{if .Spanner.IsDelete}}{{template "spanner_client_streaming_delete" .}}{{end}}
		////////////////////////////// NOTE //////////////////////////////////////
		// In the future, we might do apply if muts gets really big,  but for now,
		// we only do one apply on the database with all the records stored in muts
		//////////////////////////////////////////////////////////////////////////
		// This isn't technically an "after" hook because it happens before the mutations
		// are applied to the database.  Think of it as more of an aggregation hook for
		// gathering the proto response
		{{template "after_hook" .}}
	}
	_, err = s.SpannerDB.Apply(context.Background(), muts)
	if err != nil {
		if strings.Contains(err.Error(), "already exists") {
			return grpc.Errorf(codes.AlreadyExists, err.Error())
		} else {
			return grpc.Errorf(codes.Unknown, err.Error())
		}
	}
	stream.SendAndClose(&res)
	return nil
}
{{end}}`

const SpannerClientStreamingUpdateTemplate = `{{define "spanner_client_streaming_update"}}//spanner client streaming update
{{template "declare_spanner_arg_map" .}}

muts = append(muts, spanner.UpdateMap("{{.Spanner.TableName}}", params))
{{end}}`

const SpannerClientStreamingInsertTemplate = `{{define "spanner_client_streaming_insert"}}//spanner client streaming insert
{{template "declare_spanner_arg_slice" .}}

	muts = append(muts, spanner.Insert("{{.Spanner.TableName}}", {{.Spanner.InsertColsAsString}}, params))
{{end}}`

const SpannerClientStreamingDeleteTemplate = `{{define "spanner_client_streaming_delete"}}//spanner client streaming delete
{{template "declare_spanner_delete_key" .}}

	muts = append(muts, spanner.DeleteKeyRange({{.Spanner.TableName}}, key))
{{end}}`

const SpannerServerStreamingMethodTemplate = `{{define "spanner_server_streaming_method"}}// spanner server streaming {{.GetName}}
func (s *{{.GetServiceName}}Impl) {{.GetName}}(req *{{.GetInputType}}, stream {{.GetServiceName}}_{{.GetName}}Server) error {
	var (
	{{range $field, $type := .GetFieldsWithLocalTypesFor .GetOutputTypeStruct}}
		{{$field}} {{$type}}{{end}}
	)

	{{template "before_hook" .}}

	{{if ne (len .Spanner.QueryArgs) 0}}
	var err error
	{{end}}

	{{template "declare_spanner_arg_map" .}}

	stmt := spanner.Statement{SQL: "{{.Spanner.Query}}", Params: params}
	tx := s.SpannerDB.Single()
	defer tx.Close()
	iter := tx.Query(context.Background(), stmt)
	defer iter.Stop()
	for {
		row, err := iter.Next()
		if err == iterator.Done {
			break
		} else if err != nil {
			return grpc.Errorf(codes.Unknown, err.Error())
		}

		// scan our values out of the row
		{{range $index, $t := .GetTypeDescArrayForStruct .GetOutputTypeStruct}}
		{{if $t.IsMapped}}
		gcv := new(spanner.GenericColumnValue)
		err = row.ColumnByName("{{$t.ProtoName}}", gcv)
		if err != nil {
			return grpc.Errorf(codes.Unknown, err.Error())
		}
		err = {{$t.Name}}.SpannerScan(gcv)
		if err != nil {
			return grpc.Errorf(codes.Unknown, err.Error())
		}
		{{else}}
		err = row.ColumnByName("{{$t.ProtoName}}", &{{$t.Name}})
		if err != nil {
			return grpc.Errorf(codes.Unknown, err.Error())
		}
		{{end}}{{end}}
		res := {{.GetOutputType}}{
		{{range $field, $type := .GetTypeDescForFieldsInStruct .GetOutputTypeStruct}}
		{{$field}}: {{template "addr" $type}}{{template "base" $type}}{{template "mapping" $type}},{{end}}
		}

		{{template "after_hook" .}}

		stream.Send(&res)
	}
	return  nil
}
{{end}}`

const SpannerBidiStreamingMethodTemplate = `{{define "spanner_bidi_streaming_method"}}// spanner bidi streaming {{.GetName}}
unimplemented
{{end}}`

const SpannerHelperTemplates = `
{{define "type_desc_to_def_map"}}
{{if .IsMapped}}
	conv, err = {{.GoName}}{}.ToSpanner(req.{{.Name}}).SpannerValue()
{{else}}
	conv = req.{{.Name}}
{{end}}{{end}}

{{define "type_desc_to_def_slice"}}
{{if .IsMapped}}
	conv, err = {{.GoName}}{}.ToSpanner(req.{{.Name}}).SpannerValue()
{{else}}
	conv = req.{{.Name}}
{{end}}{{end}}

{{define "return_err_on_method"}}
	if err != nil {
{{if .IsUnary}}
		return nil, grpc.Errorf(codes.Unknown, err.Error())
{{else}}
		return grpc.Errorf(codes.Unknown, err.Error())
{{end}}
	}
{{end}}

{{define "declare_spanner_arg_map"}}
{{$method := .}}
	params := make(map[string]interface{})
{{if gt (len .Spanner.OptionArguments) 0}}
	var conv interface{}
{{end}}
{{range $key, $val := .Spanner.QueryArgs}}
{{if $val.IsFieldValue}}
	{{template "type_desc_to_def_map" $val.Field}}
	{{template "return_err_on_method" $method}}
	params["{{$val.Name}}"] = conv
{{else}}
{{if $val.IsValue}}
	conv = {{$val.Value}}
	params["{{$val.Name}}"] = conv
{{end}}
{{end}}{{end}}
{{end}}

{{define "declare_spanner_arg_slice"}}
{{$method := .}}
	params := make([]interface{}, 0)
{{if gt (len .Spanner.OptionArguments) 0}}
	var conv interface{}
{{end}}

{{range $index, $val := .Spanner.QueryArgs}}
{{if $val.IsFieldValue}}
	{{template "type_desc_to_def_slice" $val.Field}}
	{{template "return_err_on_method" $method}}
	params = append(params, conv)
{{else}}
{{if $val.IsValue}}
	params = append(params, {{$val.Value}})
{{end}}
{{end}}{{end}}
{{end}}

{{define "declare_spanner_delete_key"}}
{{$method := .}}
	start := make([]interface{}, 0)
	end := make([]interface{}, 0)
{{if gt (len .Spanner.OptionArguments) 0}}
	var conv interface{}
{{end}}
{{range $index, $arg := .Spanner.KeyRangeDesc.Start}}
{{if $arg.IsFieldValue}}
{{template "type_desc_to_def_slice" $arg.Field}}
{{template "return_err_on_method" $method}}
	start = append(start, conv)
{{else}}
	start = append(start, {{$arg.Value}})
{{end}}{{end}}
{{range $index, $arg := .Spanner.KeyRangeDesc.End}}
{{if $arg.IsFieldValue}}
{{template "type_desc_to_def_slice" $arg.Field}}
{{template "return_err_on_method" $method}}
	end = append(end, conv)
{{else}}
{{if $arg.IsValue}}
	end = append(end, {{$arg.Value}})
{{end}}
{{end}}{{end}}
	key := spanner.KeyRange{
		Start: start,
		End: end,
		Kind: spanner.{{.Spanner.KeyRangeDesc.Kind}},
	}
{{end}}

`
<|endoftext|>


<commit_before><commit_msg>Add Piece.State<commit_after><|endoftext|>


<commit_before>// Copyright 2017, TCN Inc.
// All rights reserved.

// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:

//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of TCN Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.

// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package templates

const SpannerUnaryMethodTemplate = `{{define "spanner_unary_method"}}// spanner unary select {{.GetName}}
func (s* {{.GetServiceName}}Impl) {{.GetName}} (ctx context.Context, req *{{.GetInputType}}) (*{{.GetOutputType}}, error) {
{{if .Spanner.IsSelect}}{{template "spanner_unary_select" .}}{{end}}
{{if .Spanner.IsUpdate}}{{template "spanner_unary_update" .}}{{end}}
{{if .Spanner.IsInsert}}{{template "spanner_unary_insert" .}}{{end}}
{{if .Spanner.IsDelete}}{{template "spanner_unary_delete" .}}{{end}}
{{end}}`

const SpannerHelperTemplates = `
{{define "type_desc_to_def"}}
{{if .IsMapped}}
	//is mapped
	conv, err = GoName{}.ToSpanner(req.{{.Name}}).Value()
	if err != nil {
		return nil, grpc.Errorf(codes.Unknown, err.Error())
	}
{{else}}
	// is not mapped
	conv = req.{{.Name}}
{{end}}{{end}}
`
const SpannerUnarySelectTemplate = `{{define "spanner_unary_select"}}
	var (
{{range $field, $type := .GetFieldsWithLocalTypesFor .GetOutputTypeStruct}}
		{{$field}} {{$type}}{{end}}
	)
	params := make(map[string]interface{})

	var conv string
	var err error
	//.GetSpannerSelectArgs
{{range $key, $val := .GetSpannerSelectArgs}}
{{if $val.IsFieldValue}}
	//if is.IsFieldValue
	{{template "type_desc_to_def" $val.Field}}
	params["{{$val.Name}}"] = conv
{{else}}
	//else
	//conv = { {$val.Value} }
	conv = {{$val.Value}}
	//params[{ {$val.Name} }] = conv
	params["{{$val.Name}}"] = conv
	{{end}}{{end}}

	//stmt := spanner.Statement{SQL: "{ {.Spanner.Query} }", Params: params}
	stmt := spanner.Statement{SQL: "{{.Spanner.Query}}", Params: params}
	tx := s.Client.Single()
	defer tx.Close()
	iter := tx.Query(ctx, stmt)
	rows := s.SRH.NewRowsFromIter(iter)
	rows.Next()
	if err = rows.Err(); err != nil {
		return nil, grpc.Errorf(codes.Unknown, err.Error())
	}

	//err = rows.Scan({ {range $index, $t := .GetTypeDescArrayForStruct .GetOutputTypeStruct} } &{ {$t.Name} },{ {end} })
	err = rows.Scan({{range $index, $t := .GetTypeDescArrayForStruct .GetOutputTypeStruct}} &{{$t.Name}},{{end}})
	if err == sql.ErrNoRows {
		return nil, grpc.Errorf(codes.NotFound, "%+v doesn't exist", req)
	} else if err != nil {
		return nil, grpc.Errorf(codes.Unknown, err.Error())
	}

	//res := &{ {.GetOutputType} }{
	res := &{{.GetOutputType}}{
	{{range $field, $type := .GetTypeDescForFieldsInStruct .GetOutputTypeStruct}}
	{{$field}}: {{template "addr" $type}}{{template "base" $type}}{{template "mapping" $type}},{{end}}
	}
	return res, nil
}
{{end}}`

const SpannerUnaryInsertTemplate = `{{define "spanner_unary_insert"}}
	params := []interface{}{
		{{range $index, $val := .GetSpannerInsertArgs "req"}}
		{{$val}},{{end}}
	}
	muts := make([]*spanner.Mutation, 1)
	muts[0] = spanner.Insert("{{.Spanner.TableName}}", {{.Spanner.InsertCols}}, params)
	_, err := s.SpannerDB.Apply(muts)
	if err != nil {
		if strings.Contains(err.Error(), "already exists") {
			return grpc.Errorf(codes.AlreadyExists, err.Error())
		} else {
			returnn nil, grpc.Errorf(codes.Unknown, err.Error())
		}
	}
	res := &{{.GetOutputType}}{}

	return res, nil
}
{{end}}`

const SpannerUnaryUpdateTemplate = `{{define "spanner_unary_update"}}
	params := map[string]interface{}{
		{{range $key, $val := .GetSpannerUpdateArgs "req"}}
		{{$key}}: {{$val}},\n{{end}}
	}
	muts := make([]*spanner.Mutation, 1)
	muts[0] = spanner.UpdateMap("{{.Spanner.TableName}}", params)
	_, err := s.SpannerDB.Apply(muts)
	if err != nil {
		if strings.Contains(err.Error(), "already exists") {
			return grpc.Errorf(codes.AlreadyExists, err.Error())
		} else {
			return nil, grpc.Errorf(codes.Unknown, err.Error())
		}
	}
	res := &{{.GetOutputType}}{}

	return res, nil
}
{{end}}`

const SpannerUnaryDeleteTemplate = `{{define "spanner_unary_delete"}}
	key := {{.GetDeleteKeyRange}}
	muts := make([]*spanner.Mutation, 1)
	muts[0] = spanner.DeleteKeyRange("{{.Spanner.TableName}}", key)
	_, err := s.SpannerDB.Apply(muts)
	if err != nil {
		if strings.Contains(err.Error(), "does not exist") {
			return grpc.Errorf(codes.NotFound, err.Error())
		}
	}
{{end}}`

const SpannerClientStreamingMethodTemplate = `{{define "spanner_client_streaming_method"}}// spanner client streaming {{.GetName}}
func (s *{{.GetServiceName}}Impl) {{.GetName}}(stream {{.GetServiceName}}_{{.GetName}}Server) error {
	var totalAffected int64
	muts := make([]*spanner.Mutation, 0)
	for {
		req, err := stream.Recv()
		if err == io.EOF {
			break
		} else if err != nil {
			return grpc.Errorf(codes.Unknown, err.Error())
		}
		totalAffected += 1

		{{if .Spanner.IsInsert}}{{template "spanner_client_streaming_insert" .}}{{end}}
		{{if .Spanner.IsUpdate}}{{template "spanner_client_streaming_update" .}}{{end}}
		{{if .Spanner.IsDelete}}{{template "spanner_client_streaming_delete" .}}{{end}}
		//In the future, we might do apply if muts gets really big,  but for now,
		// we only do one apply on the database with all the records stored in muts
	}
	_, err := s.SpannerDB.Apply(muts)
	if err != nil {
		if strings.Contains(err.Error(), "already exists") {
			return grpc.Errorf(codes.AlreadyExists, err.Error())
		} else {
			return nil, grpc.Errorf(codes.Unknown, err.Error())
		}
	}
	stream.SendAndClose(&{{.GetOutputType}}{Count: totalAffeted})

{{end}}`

const SpannerClientStreamingUpdateTemplate = `{{define "spanner_client_streaming_update"}}//spanner client streaming update
params := map[string]interface{}{
{{range $key, $val := .GetSpannerUpdateArgs}}
	{{$key}}: {{$val}},\n{{end}}
}
muts = append(muts, spanner.UpdateMap("{{.Spanner.TableName}}", params))
{{end}}`

const SpannerClientStreamingInsertTemplate = `{{define "spanner_client_streaming_insert"}}//spanner client streaming update
params := []interface{}{
{{range $index, $val := .GetSpannerInsertArgs}}
	{{$val}},\n{{end}}
}
muts = append(muts, spanner.Insert("{{.Spanner.TableName}}", {{.Spanner.InsertColsAsString}}, params))
{{end}}`

const SpannerClientStreamingDeleteTemplate = `{{define "spanner_client_streaming_delete"}}//spanner client streaming update
key := {{.GetDeleteKeyRange "req"}}
muts = append(muts, spanner.DeleteKeyRange("{{.Spanner.TableName}}", key)
{{end}}`

const SpannerServerStreamingMethodTemplate = `{{define "spanner_server_streaming_method"}}// spanner server streaming {{.GetName}}
func (s *{{.GetServiceName}}Impl) {{.GetName}}(req *{{.GetInputType}}, stream {{.GetServiceName}}_{{.GetName}}Server) error {
	var (
	{{range $field, $type := .GetFieldsWithLocalTypesFor .GetOutputTypeStruct}}
		{{$field}} {{$type}}{{end}}
	)
	params := make(map[string]interface{})

	var conv string
	var err error
	//.GetSpannerSelectArgs
{{range $key, $val := .GetSpannerSelectArgs}}
{{if $val.IsFieldValue}}
	//if is.IsFieldValue
	{{template "type_desc_to_def" $val.Field}}
	params["{{$val.Name}}"] = conv
{{else}}
	//else
	//conv = { {$val.Value} }
	conv = {{$val.Value}}
	//params["{ {$val.Name} }"] = conv
	params["{{$val.Name}}"] = conv
{{end}}{{end}}

	stmt := spanner.Statement{SQL: "{{.Spanner.Query}}", Params: params}
	tx := s.Client.Single()
	defer tx.Close()
	iter := tx.Query(ctx, stmt)
	rows := s.SRH.NewRowsFromIter(iter)
	for rows.Next() {
		if err := rows.Err(); err != nil {
			if err == sql.ErrNowRows {
				return grpc.Errorf(codes.NotFound, "%+v doesn't exist", req)
			}
			return grpc.Errorf(codes.nknown, err.Error())
		}
		err := rows.Scan({{range $index, $t := .GetTypeDescArrayForStruct .GetOutputTypeStruct}} &{{$t.Name}},{{end}})
		if err != nil {
			return grpc.Errorf(codes.Unknown, err.Error())
		}
		res := &{{.GetOutputType}}{
		{{range $field, $type := .GetTypeDescForFieldsInStruct .GetOutputTypeStruct}}
		{{$field}}: {{template "addr" $type}}{{template "base" $type}}{{template "mapping" $type}},{{end}}
		}
		stream.Send(res)
	}
	return  nil
}
{{end}}`

const SpannerBidiStreamingMethodTemplate = `{{define "spanner_bidi_streaming_method"}}// spanner bidi streaming {{.GetName}} unimplemented{{end}}`
<commit_msg>fixed client stream template<commit_after>// Copyright 2017, TCN Inc.
// All rights reserved.

// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:

//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of TCN Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.

// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package templates

const SpannerUnaryMethodTemplate = `{{define "spanner_unary_method"}}// spanner unary select {{.GetName}}
func (s* {{.GetServiceName}}Impl) {{.GetName}} (ctx context.Context, req *{{.GetInputType}}) (*{{.GetOutputType}}, error) {
{{if .Spanner.IsSelect}}{{template "spanner_unary_select" .}}{{end}}
{{if .Spanner.IsUpdate}}{{template "spanner_unary_update" .}}{{end}}
{{if .Spanner.IsInsert}}{{template "spanner_unary_insert" .}}{{end}}
{{if .Spanner.IsDelete}}{{template "spanner_unary_delete" .}}{{end}}
{{end}}`

const SpannerHelperTemplates = `
{{define "type_desc_to_def"}}
{{if .IsMapped}}
	//is mapped
	conv, err = GoName{}.ToSpanner(req.{{.Name}}).Value()
	if err != nil {
		return nil, grpc.Errorf(codes.Unknown, err.Error())
	}
{{else}}
	// is not mapped
	conv = req.{{.Name}}
{{end}}{{end}}
`
const SpannerUnarySelectTemplate = `{{define "spanner_unary_select"}}
	var (
{{range $field, $type := .GetFieldsWithLocalTypesFor .GetOutputTypeStruct}}
		{{$field}} {{$type}}{{end}}
	)
	params := make(map[string]interface{})

	var conv string
	var err error
	//.GetSpannerSelectArgs
{{range $key, $val := .GetSpannerSelectArgs}}
{{if $val.IsFieldValue}}
	//if is.IsFieldValue
	{{template "type_desc_to_def" $val.Field}}
	params["{{$val.Name}}"] = conv
{{else}}
	//else
	//conv = { {$val.Value} }
	conv = {{$val.Value}}
	//params[{ {$val.Name} }] = conv
	params["{{$val.Name}}"] = conv
	{{end}}{{end}}

	//stmt := spanner.Statement{SQL: "{ {.Spanner.Query} }", Params: params}
	stmt := spanner.Statement{SQL: "{{.Spanner.Query}}", Params: params}
	tx := s.Client.Single()
	defer tx.Close()
	iter := tx.Query(ctx, stmt)
	rows := s.SRH.NewRowsFromIter(iter)
	rows.Next()
	if err = rows.Err(); err != nil {
		return nil, grpc.Errorf(codes.Unknown, err.Error())
	}

	//err = rows.Scan({ {range $index, $t := .GetTypeDescArrayForStruct .GetOutputTypeStruct} } &{ {$t.Name} },{ {end} })
	err = rows.Scan({{range $index, $t := .GetTypeDescArrayForStruct .GetOutputTypeStruct}} &{{$t.Name}},{{end}})
	if err == sql.ErrNoRows {
		return nil, grpc.Errorf(codes.NotFound, "%+v doesn't exist", req)
	} else if err != nil {
		return nil, grpc.Errorf(codes.Unknown, err.Error())
	}

	//res := &{ {.GetOutputType} }{
	res := &{{.GetOutputType}}{
	{{range $field, $type := .GetTypeDescForFieldsInStruct .GetOutputTypeStruct}}
	{{$field}}: {{template "addr" $type}}{{template "base" $type}}{{template "mapping" $type}},{{end}}
	}
	return res, nil
}
{{end}}`

const SpannerUnaryInsertTemplate = `{{define "spanner_unary_insert"}}
	params := []interface{}{
		{{range $index, $val := .GetSpannerInsertArgs "req"}}
		{{$val}},{{end}}
	}
	muts := make([]*spanner.Mutation, 1)
	muts[0] = spanner.Insert("{{.Spanner.TableName}}", {{.Spanner.InsertCols}}, params)
	_, err := s.SpannerDB.Apply(muts)
	if err != nil {
		if strings.Contains(err.Error(), "already exists") {
			return grpc.Errorf(codes.AlreadyExists, err.Error())
		} else {
			returnn nil, grpc.Errorf(codes.Unknown, err.Error())
		}
	}
	res := &{{.GetOutputType}}{}

	return res, nil
}
{{end}}`

const SpannerUnaryUpdateTemplate = `{{define "spanner_unary_update"}}
	params := map[string]interface{}{
		{{range $key, $val := .GetSpannerUpdateArgs "req"}}
		{{$key}}: {{$val}},\n{{end}}
	}
	muts := make([]*spanner.Mutation, 1)
	muts[0] = spanner.UpdateMap("{{.Spanner.TableName}}", params)
	_, err := s.SpannerDB.Apply(muts)
	if err != nil {
		if strings.Contains(err.Error(), "already exists") {
			return grpc.Errorf(codes.AlreadyExists, err.Error())
		} else {
			return nil, grpc.Errorf(codes.Unknown, err.Error())
		}
	}
	res := &{{.GetOutputType}}{}

	return res, nil
}
{{end}}`

const SpannerUnaryDeleteTemplate = `{{define "spanner_unary_delete"}}
	key := {{.GetDeleteKeyRange}}
	muts := make([]*spanner.Mutation, 1)
	muts[0] = spanner.DeleteKeyRange("{{.Spanner.TableName}}", key)
	_, err := s.SpannerDB.Apply(muts)
	if err != nil {
		if strings.Contains(err.Error(), "does not exist") {
			return grpc.Errorf(codes.NotFound, err.Error())
		}
	}
{{end}}`

const SpannerClientStreamingMethodTemplate = `{{define "spanner_client_streaming_method"}}// spanner client streaming {{.GetName}}
func (s *{{.GetServiceName}}Impl) {{.GetName}}(stream {{.GetServiceName}}_{{.GetName}}Server) error {
	var totalAffected int64
	muts := make([]*spanner.Mutation, 0)
	for {
		req, err := stream.Recv()
		if err == io.EOF {
			break
		} else if err != nil {
			return grpc.Errorf(codes.Unknown, err.Error())
		}
		totalAffected += 1

		{{if .Spanner.IsInsert}}{{template "spanner_client_streaming_insert" .}}{{end}}
		{{if .Spanner.IsUpdate}}{{template "spanner_client_streaming_update" .}}{{end}}
		{{if .Spanner.IsDelete}}{{template "spanner_client_streaming_delete" .}}{{end}}
		//In the future, we might do apply if muts gets really big,  but for now,
		// we only do one apply on the database with all the records stored in muts
	}
	_, err := s.SpannerDB.Apply(muts)
	if err != nil {
		if strings.Contains(err.Error(), "already exists") {
			return grpc.Errorf(codes.AlreadyExists, err.Error())
		} else {
			return nil, grpc.Errorf(codes.Unknown, err.Error())
		}
	}
	stream.SendAndClose(&{{.GetOutputType}}{Count: totalAffeted})
}
{{end}}`

const SpannerClientStreamingUpdateTemplate = `{{define "spanner_client_streaming_update"}}//spanner client streaming update
params := map[string]interface{}{
{{range $key, $val := .GetSpannerUpdateArgs}}
	{{$key}}: {{$val}},\n{{end}}
}
muts = append(muts, spanner.UpdateMap("{{.Spanner.TableName}}", params))
{{end}}`

const SpannerClientStreamingInsertTemplate = `{{define "spanner_client_streaming_insert"}}//spanner client streaming update
params := []interface{}{
{{range $index, $val := .GetSpannerInsertArgs}}
	{{$val}},\n{{end}}
}
muts = append(muts, spanner.Insert("{{.Spanner.TableName}}", {{.Spanner.InsertColsAsString}}, params))
{{end}}`

const SpannerClientStreamingDeleteTemplate = `{{define "spanner_client_streaming_delete"}}//spanner client streaming update
key := {{.GetDeleteKeyRange "req"}}
muts = append(muts, spanner.DeleteKeyRange("{{.Spanner.TableName}}", key)
{{end}}`

const SpannerServerStreamingMethodTemplate = `{{define "spanner_server_streaming_method"}}// spanner server streaming {{.GetName}}
func (s *{{.GetServiceName}}Impl) {{.GetName}}(req *{{.GetInputType}}, stream {{.GetServiceName}}_{{.GetName}}Server) error {
	var (
	{{range $field, $type := .GetFieldsWithLocalTypesFor .GetOutputTypeStruct}}
		{{$field}} {{$type}}{{end}}
	)
	params := make(map[string]interface{})

	var conv string
	var err error
	//.GetSpannerSelectArgs
{{range $key, $val := .GetSpannerSelectArgs}}
{{if $val.IsFieldValue}}
	//if is.IsFieldValue
	{{template "type_desc_to_def" $val.Field}}
	params["{{$val.Name}}"] = conv
{{else}}
	//else
	//conv = { {$val.Value} }
	conv = {{$val.Value}}
	//params["{ {$val.Name} }"] = conv
	params["{{$val.Name}}"] = conv
{{end}}{{end}}

	stmt := spanner.Statement{SQL: "{{.Spanner.Query}}", Params: params}
	tx := s.Client.Single()
	defer tx.Close()
	iter := tx.Query(ctx, stmt)
	rows := s.SRH.NewRowsFromIter(iter)
	for rows.Next() {
		if err := rows.Err(); err != nil {
			if err == sql.ErrNowRows {
				return grpc.Errorf(codes.NotFound, "%+v doesn't exist", req)
			}
			return grpc.Errorf(codes.nknown, err.Error())
		}
		err := rows.Scan({{range $index, $t := .GetTypeDescArrayForStruct .GetOutputTypeStruct}} &{{$t.Name}},{{end}})
		if err != nil {
			return grpc.Errorf(codes.Unknown, err.Error())
		}
		res := &{{.GetOutputType}}{
		{{range $field, $type := .GetTypeDescForFieldsInStruct .GetOutputTypeStruct}}
		{{$field}}: {{template "addr" $type}}{{template "base" $type}}{{template "mapping" $type}},{{end}}
		}
		stream.Send(res)
	}
	return  nil
}
{{end}}`

const SpannerBidiStreamingMethodTemplate = `{{define "spanner_bidi_streaming_method"}}// spanner bidi streaming {{.GetName}} unimplemented{{end}}`
<|endoftext|>


<commit_before><commit_msg>Create pipes.go<commit_after><|endoftext|>


<commit_before>package rpcserver

import (
	"context"
	"fmt"
	"github.com/getlantern/bytemap"
	"github.com/getlantern/errors"
	"github.com/getlantern/golog"
	"github.com/getlantern/wal"
	"github.com/getlantern/zenodb"
	"github.com/getlantern/zenodb/common"
	"github.com/getlantern/zenodb/core"
	"github.com/getlantern/zenodb/encoding"
	"github.com/getlantern/zenodb/planner"
	"github.com/getlantern/zenodb/rpc"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
	"net"
	"time"
)

var (
	log = golog.LoggerFor("zenodb.rpc")
)

type Opts struct {
	// Password, if specified, is the password that clients must present in order
	// to access the server.
	Password string
}

// DB is an interface for database-like things (implemented by common.DB).
type DB interface {
	InsertRaw(stream string, ts time.Time, dims bytemap.ByteMap, vals bytemap.ByteMap) error

	Query(sqlString string, isSubQuery bool, subQueryResults [][]interface{}, includeMemStore bool) (core.FlatRowSource, error)

	Follow(f *common.Follow, cb func([]byte, wal.Offset) error)

	RegisterQueryHandler(partition int, query planner.QueryClusterFN)
}

func Serve(db DB, l net.Listener, opts *Opts) error {
	l = &rpc.SnappyListener{l}
	gs := grpc.NewServer(grpc.CustomCodec(rpc.Codec))
	gs.RegisterService(&rpc.ServiceDesc, &server{db, opts.Password})
	return gs.Serve(l)
}

type server struct {
	db       DB
	password string
}

func (s *server) Insert(stream grpc.ServerStream) error {
	// No need to authorize, anyone can insert

	now := time.Now()
	streamName := ""

	report := &rpc.InsertReport{
		Errors: make(map[int]string),
	}

	i := -1
	for {
		i++
		insert := &rpc.Insert{}
		err := stream.RecvMsg(insert)
		if err != nil {
			return fmt.Errorf("Error reading insert: %v", err)
		}
		if insert.EndOfInserts {
			// We're done inserting
			return stream.SendMsg(report)
		}
		report.Received++

		if streamName == "" {
			streamName = insert.Stream
			if streamName == "" {
				return fmt.Errorf("Please specify a stream")
			}
		}

		if len(insert.Dims) == 0 {
			report.Errors[i] = fmt.Sprintf("Need at least one dim")
			continue
		}
		if len(insert.Vals) == 0 {
			report.Errors[i] = fmt.Sprintf("Need at least one val")
			continue
		}
		var ts time.Time
		if insert.TS == 0 {
			ts = now
		} else {
			ts = encoding.TimeFromInt(insert.TS)
		}

		// TODO: make sure we don't barf on invalid bytemaps here
		insertErr := s.db.InsertRaw(streamName, ts, bytemap.ByteMap(insert.Dims), bytemap.ByteMap(insert.Vals))
		if insertErr != nil {
			report.Errors[i] = fmt.Sprintf("Unable to insert: %v", insertErr)
			continue
		}
		report.Succeeded++
	}
}

func (s *server) Query(q *rpc.Query, stream grpc.ServerStream) error {
	authorizeErr := s.authorize(stream)
	if authorizeErr != nil {
		return authorizeErr
	}

	source, err := s.db.Query(q.SQLString, q.IsSubQuery, q.SubQueryResults, q.IncludeMemStore)
	if err != nil {
		return err
	}

	rr := &rpc.RemoteQueryResult{}
	stats, err := source.Iterate(stream.Context(), func(fields core.Fields) error {
		// Send query metadata
		md := zenodb.MetaDataFor(source, fields)
		return stream.SendMsg(md)
	}, func(row *core.FlatRow) (bool, error) {
		rr.Row = row
		return true, stream.SendMsg(rr)
	})
	if err != nil {
		return err
	}

	// Send end of results
	rr.Row = nil
	if stats != nil {
		rr.Stats = stats.(*common.QueryStats)
	}
	rr.EndOfResults = true
	return stream.SendMsg(rr)
}

func (s *server) Follow(f *common.Follow, stream grpc.ServerStream) error {
	authorizeErr := s.authorize(stream)
	if authorizeErr != nil {
		return authorizeErr
	}

	log.Debugf("Follower %d joined", f.PartitionNumber)
	defer log.Debugf("Follower %d left", f.PartitionNumber)
	s.db.Follow(f, func(data []byte, newOffset wal.Offset) error {
		return stream.SendMsg(&rpc.Point{data, newOffset})
	})
	return nil
}

func (s *server) HandleRemoteQueries(r *rpc.RegisterQueryHandler, stream grpc.ServerStream) error {
	initialResultCh := make(chan *rpc.RemoteQueryResult)
	initialErrCh := make(chan error, 1)
	finalErrCh := make(chan error, 1)

	finish := func(err error) {
		select {
		case finalErrCh <- err:
			// ok
			log.Debugf("Posted final err for partition %d: %v", r.Partition, err)
		default:
			// ignore
		}
	}

	s.db.RegisterQueryHandler(r.Partition, func(ctx context.Context, sqlString string, isSubQuery bool, subQueryResults [][]interface{}, unflat bool, onFields core.OnFields, onRow core.OnRow, onFlatRow core.OnFlatRow) (interface{}, error) {
		q := &rpc.Query{
			SQLString:       sqlString,
			IsSubQuery:      isSubQuery,
			SubQueryResults: subQueryResults,
			Unflat:          unflat,
			IncludeMemStore: common.ShouldIncludeMemStore(ctx),
		}
		q.Deadline, q.HasDeadline = ctx.Deadline()
		sendErr := stream.SendMsg(q)

		m, recvErr := <-initialResultCh, <-initialErrCh

		// Check send error after reading initial result to avoid blocking
		// unnecessarily
		if sendErr != nil {
			err := errors.New("Unable to send query: %v", sendErr)
			finish(err)
			return nil, err
		}

		var finalErr error

		first := true
	receiveLoop:
		for {
			// Process current result
			if recvErr != nil {
				m.Error = recvErr.Error()
				finalErr = errors.New("Unable to receive result: %v", recvErr)
				break
			}

			log.Debugf("For partition %d, result has error? %v", r.Partition, m.Error)
			if first {
				// First message contains only fields information
				onFields(m.Fields)
				first = false
			} else {
				// Subsequent messages contain data
				if m.EndOfResults {
					break
				}
				var more bool
				var err error
				if unflat {
					more, err = onRow(m.Key, m.Vals)
				} else {
					more, err = onFlatRow(m.Row)
				}
				if !more || err != nil {
					log.Debugf("Setting final error for partition %d: %v", r.Partition, err)
					finalErr = err
					break receiveLoop
				}
			}

			// Read next result
			m = &rpc.RemoteQueryResult{}
			recvErr = stream.RecvMsg(m)
		}

		log.Debugf("Finishing partition %d with finalErr: %v", finalErr)
		finish(finalErr)
		return m.Stats, finalErr
	})

	// Block on reading initial result to keep connection open
	m := &rpc.RemoteQueryResult{}
	err := stream.RecvMsg(m)
	initialResultCh <- m
	initialErrCh <- err

	if err == nil {
		// Wait for final error so we don't close the connection prematurely
		err = <-finalErrCh
	}
	log.Debugf("Returning err for partition %d?: %v", r.Partition, err)
	return err
}

func (s *server) authorize(stream grpc.ServerStream) error {
	if s.password == "" {
		log.Debug("No password specified, allowing access to world")
		return nil
	}
	md, ok := metadata.FromIncomingContext(stream.Context())
	if !ok {
		return log.Error("No metadata provided, unable to authenticate")
	}
	passwords := md[rpc.PasswordKey]
	for _, password := range passwords {
		if password == s.password {
			// authorized
			return nil
		}
	}
	return log.Error("None of the provided passwords matched, not authorized!")
}
<commit_msg>Added more logging for query error handling<commit_after>package rpcserver

import (
	"context"
	"fmt"
	"github.com/getlantern/bytemap"
	"github.com/getlantern/errors"
	"github.com/getlantern/golog"
	"github.com/getlantern/wal"
	"github.com/getlantern/zenodb"
	"github.com/getlantern/zenodb/common"
	"github.com/getlantern/zenodb/core"
	"github.com/getlantern/zenodb/encoding"
	"github.com/getlantern/zenodb/planner"
	"github.com/getlantern/zenodb/rpc"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
	"net"
	"time"
)

var (
	log = golog.LoggerFor("zenodb.rpc")
)

type Opts struct {
	// Password, if specified, is the password that clients must present in order
	// to access the server.
	Password string
}

// DB is an interface for database-like things (implemented by common.DB).
type DB interface {
	InsertRaw(stream string, ts time.Time, dims bytemap.ByteMap, vals bytemap.ByteMap) error

	Query(sqlString string, isSubQuery bool, subQueryResults [][]interface{}, includeMemStore bool) (core.FlatRowSource, error)

	Follow(f *common.Follow, cb func([]byte, wal.Offset) error)

	RegisterQueryHandler(partition int, query planner.QueryClusterFN)
}

func Serve(db DB, l net.Listener, opts *Opts) error {
	l = &rpc.SnappyListener{l}
	gs := grpc.NewServer(grpc.CustomCodec(rpc.Codec))
	gs.RegisterService(&rpc.ServiceDesc, &server{db, opts.Password})
	return gs.Serve(l)
}

type server struct {
	db       DB
	password string
}

func (s *server) Insert(stream grpc.ServerStream) error {
	// No need to authorize, anyone can insert

	now := time.Now()
	streamName := ""

	report := &rpc.InsertReport{
		Errors: make(map[int]string),
	}

	i := -1
	for {
		i++
		insert := &rpc.Insert{}
		err := stream.RecvMsg(insert)
		if err != nil {
			return fmt.Errorf("Error reading insert: %v", err)
		}
		if insert.EndOfInserts {
			// We're done inserting
			return stream.SendMsg(report)
		}
		report.Received++

		if streamName == "" {
			streamName = insert.Stream
			if streamName == "" {
				return fmt.Errorf("Please specify a stream")
			}
		}

		if len(insert.Dims) == 0 {
			report.Errors[i] = fmt.Sprintf("Need at least one dim")
			continue
		}
		if len(insert.Vals) == 0 {
			report.Errors[i] = fmt.Sprintf("Need at least one val")
			continue
		}
		var ts time.Time
		if insert.TS == 0 {
			ts = now
		} else {
			ts = encoding.TimeFromInt(insert.TS)
		}

		// TODO: make sure we don't barf on invalid bytemaps here
		insertErr := s.db.InsertRaw(streamName, ts, bytemap.ByteMap(insert.Dims), bytemap.ByteMap(insert.Vals))
		if insertErr != nil {
			report.Errors[i] = fmt.Sprintf("Unable to insert: %v", insertErr)
			continue
		}
		report.Succeeded++
	}
}

func (s *server) Query(q *rpc.Query, stream grpc.ServerStream) error {
	authorizeErr := s.authorize(stream)
	if authorizeErr != nil {
		return authorizeErr
	}

	source, err := s.db.Query(q.SQLString, q.IsSubQuery, q.SubQueryResults, q.IncludeMemStore)
	if err != nil {
		return err
	}

	rr := &rpc.RemoteQueryResult{}
	stats, err := source.Iterate(stream.Context(), func(fields core.Fields) error {
		// Send query metadata
		md := zenodb.MetaDataFor(source, fields)
		return stream.SendMsg(md)
	}, func(row *core.FlatRow) (bool, error) {
		rr.Row = row
		return true, stream.SendMsg(rr)
	})
	if err != nil {
		return err
	}

	// Send end of results
	rr.Row = nil
	if stats != nil {
		rr.Stats = stats.(*common.QueryStats)
	}
	rr.EndOfResults = true
	return stream.SendMsg(rr)
}

func (s *server) Follow(f *common.Follow, stream grpc.ServerStream) error {
	authorizeErr := s.authorize(stream)
	if authorizeErr != nil {
		return authorizeErr
	}

	log.Debugf("Follower %d joined", f.PartitionNumber)
	defer log.Debugf("Follower %d left", f.PartitionNumber)
	s.db.Follow(f, func(data []byte, newOffset wal.Offset) error {
		return stream.SendMsg(&rpc.Point{data, newOffset})
	})
	return nil
}

func (s *server) HandleRemoteQueries(r *rpc.RegisterQueryHandler, stream grpc.ServerStream) error {
	initialResultCh := make(chan *rpc.RemoteQueryResult)
	initialErrCh := make(chan error, 1)
	finalErrCh := make(chan error, 1)

	finish := func(err error) {
		select {
		case finalErrCh <- err:
			// ok
			log.Debugf("Posted final err for partition %d: %v", r.Partition, err)
		default:
			// ignore
		}
	}

	s.db.RegisterQueryHandler(r.Partition, func(ctx context.Context, sqlString string, isSubQuery bool, subQueryResults [][]interface{}, unflat bool, onFields core.OnFields, onRow core.OnRow, onFlatRow core.OnFlatRow) (interface{}, error) {
		q := &rpc.Query{
			SQLString:       sqlString,
			IsSubQuery:      isSubQuery,
			SubQueryResults: subQueryResults,
			Unflat:          unflat,
			IncludeMemStore: common.ShouldIncludeMemStore(ctx),
		}
		q.Deadline, q.HasDeadline = ctx.Deadline()
		sendErr := stream.SendMsg(q)

		m, recvErr := <-initialResultCh, <-initialErrCh

		// Check send error after reading initial result to avoid blocking
		// unnecessarily
		if sendErr != nil {
			err := errors.New("Unable to send query: %v", sendErr)
			finish(err)
			return nil, err
		}

		var finalErr error

		first := true
	receiveLoop:
		for {
			// Process current result
			if recvErr != nil {
				m.Error = recvErr.Error()
				finalErr = errors.New("Unable to receive result: %v", recvErr)
				break
			}

			log.Debugf("For partition %d, result has error? %v", r.Partition, m.Error)
			if first {
				// First message contains only fields information
				onFields(m.Fields)
				first = false
			} else {
				if m.Error != "" {
					log.Debugf("Setting final error for partition %d: %v", r.Partition, m.Error)
					finalErr = errors.New(m.Error)
				}
				// Subsequent messages contain data
				if m.EndOfResults {
					break
				}
				var more bool
				var err error
				if unflat {
					more, err = onRow(m.Key, m.Vals)
				} else {
					more, err = onFlatRow(m.Row)
				}
				if !more || err != nil {
					log.Debugf("Setting final error for partition %d: %v", r.Partition, err)
					finalErr = err
					break receiveLoop
				}
			}

			// Read next result
			m = &rpc.RemoteQueryResult{}
			recvErr = stream.RecvMsg(m)
		}

		log.Debugf("Finishing partition %d with finalErr: %v", r.Partition, finalErr)
		finish(finalErr)
		return m.Stats, finalErr
	})

	// Block on reading initial result to keep connection open
	m := &rpc.RemoteQueryResult{}
	err := stream.RecvMsg(m)
	initialResultCh <- m
	initialErrCh <- err

	if err == nil {
		// Wait for final error so we don't close the connection prematurely
		err = <-finalErrCh
	}
	log.Debugf("Returning err for partition %d?: %v", r.Partition, err)
	return err
}

func (s *server) authorize(stream grpc.ServerStream) error {
	if s.password == "" {
		log.Debug("No password specified, allowing access to world")
		return nil
	}
	md, ok := metadata.FromIncomingContext(stream.Context())
	if !ok {
		return log.Error("No metadata provided, unable to authenticate")
	}
	passwords := md[rpc.PasswordKey]
	for _, password := range passwords {
		if password == s.password {
			// authorized
			return nil
		}
	}
	return log.Error("None of the provided passwords matched, not authorized!")
}
<|endoftext|>


<commit_before>// Copyright 2013 Prometheus Team
// 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 ast

import (
	"flag"
	"github.com/prometheus/prometheus/model"
	"github.com/prometheus/prometheus/storage/metric"
	"time"
)

var defaultStalenessDelta = flag.Int("defaultStalenessDelta", 300, "Default staleness delta allowance in seconds during expression evaluations.")

// Describes the lenience limits to apply to values from the materialized view.
type StalenessPolicy struct {
	// Describes the inclusive limit at which individual points if requested will
	// be matched and subject to interpolation.
	DeltaAllowance time.Duration
}

type viewAdapter struct {
	// Policy that dictates when sample values around an evaluation time are to
	// be interpreted as stale.
	stalenessPolicy StalenessPolicy
	// AST-global storage to use for operations that are not supported by views
	// (i.e. fingerprint->metric lookups).
	storage *metric.TieredStorage
	// The materialized view which contains all timeseries data required for
	// executing a query.
	view metric.View
}

// interpolateSamples interpolates a value at a target time between two
// provided sample pairs.
func interpolateSamples(first, second *model.SamplePair, timestamp time.Time) *model.SamplePair {
	dv := second.Value - first.Value
	dt := second.Timestamp.Sub(first.Timestamp)

	dDt := dv / model.SampleValue(dt)
	offset := model.SampleValue(timestamp.Sub(first.Timestamp))

	return &model.SamplePair{
		Value:     first.Value + (offset * dDt),
		Timestamp: timestamp,
	}
}

// chooseClosestSample chooses the closest sample of a list of samples
// surrounding a given target time. If samples are found both before and after
// the target time, the sample value is interpolated between these. Otherwise,
// the single closest sample is returned verbatim.
func (v *viewAdapter) chooseClosestSample(samples model.Values, timestamp time.Time) *model.SamplePair {
	var closestBefore *model.SamplePair
	var closestAfter *model.SamplePair
	for _, candidate := range samples {
		delta := candidate.Timestamp.Sub(timestamp)
		// Samples before target time.
		if delta < 0 {
			// Ignore samples outside of staleness policy window.
			if -delta > v.stalenessPolicy.DeltaAllowance {
				continue
			}
			// Ignore samples that are farther away than what we've seen before.
			if closestBefore != nil && candidate.Timestamp.Before(closestBefore.Timestamp) {
				continue
			}
			sample := candidate
			closestBefore = &sample
		}

		// Samples after target time.
		if delta >= 0 {
			// Ignore samples outside of staleness policy window.
			if delta > v.stalenessPolicy.DeltaAllowance {
				continue
			}
			// Ignore samples that are farther away than samples we've seen before.
			if closestAfter != nil && candidate.Timestamp.After(closestAfter.Timestamp) {
				continue
			}
			sample := candidate
			closestAfter = &sample
		}
	}

	switch {
	case closestBefore != nil && closestAfter != nil:
		return interpolateSamples(closestBefore, closestAfter, timestamp)
	case closestBefore != nil:
		return closestBefore
	default:
		return closestAfter
	}
}

func (v *viewAdapter) GetValueAtTime(fingerprints model.Fingerprints, timestamp time.Time) (samples Vector, err error) {
	for _, fingerprint := range fingerprints {
		sampleCandidates := v.view.GetValueAtTime(fingerprint, timestamp)
		samplePair := v.chooseClosestSample(sampleCandidates, timestamp)
		m, err := v.storage.GetMetricForFingerprint(fingerprint)
		if err != nil {
			continue
		}
		if samplePair != nil {
			samples = append(samples, model.Sample{
				Metric:    m,
				Value:     samplePair.Value,
				Timestamp: timestamp,
			})
		}
	}
	return samples, err
}

func (v *viewAdapter) GetBoundaryValues(fingerprints model.Fingerprints, interval *model.Interval) (sampleSets []model.SampleSet, err error) {
	for _, fingerprint := range fingerprints {
		samplePairs := v.view.GetBoundaryValues(fingerprint, *interval)
		if samplePairs == nil {
			continue
		}

		// TODO: memoize/cache this.
		m, err := v.storage.GetMetricForFingerprint(fingerprint)
		if err != nil {
			continue
		}

		sampleSet := model.SampleSet{
			Metric: m,
			Values: samplePairs,
		}
		sampleSets = append(sampleSets, sampleSet)
	}
	return sampleSets, nil
}

func (v *viewAdapter) GetRangeValues(fingerprints model.Fingerprints, interval *model.Interval) (sampleSets []model.SampleSet, err error) {
	for _, fingerprint := range fingerprints {
		samplePairs := v.view.GetRangeValues(fingerprint, *interval)
		if samplePairs == nil {
			continue
		}

		// TODO: memoize/cache this.
		m, err := v.storage.GetMetricForFingerprint(fingerprint)
		if err != nil {
			continue
		}

		sampleSet := model.SampleSet{
			Metric: m,
			Values: samplePairs,
		}
		sampleSets = append(sampleSets, sampleSet)
	}
	return sampleSets, nil
}

func NewViewAdapter(view metric.View, storage *metric.TieredStorage) *viewAdapter {
	stalenessPolicy := StalenessPolicy{
		DeltaAllowance: time.Duration(*defaultStalenessDelta) * time.Second,
	}

	return &viewAdapter{
		stalenessPolicy: stalenessPolicy,
		storage:         storage,
		view:            view,
	}
}
<commit_msg>Skip any empty range/boundary elements, not only nil ones.<commit_after>// Copyright 2013 Prometheus Team
// 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 ast

import (
	"flag"
	"github.com/prometheus/prometheus/model"
	"github.com/prometheus/prometheus/storage/metric"
	"time"
)

var defaultStalenessDelta = flag.Int("defaultStalenessDelta", 300, "Default staleness delta allowance in seconds during expression evaluations.")

// Describes the lenience limits to apply to values from the materialized view.
type StalenessPolicy struct {
	// Describes the inclusive limit at which individual points if requested will
	// be matched and subject to interpolation.
	DeltaAllowance time.Duration
}

type viewAdapter struct {
	// Policy that dictates when sample values around an evaluation time are to
	// be interpreted as stale.
	stalenessPolicy StalenessPolicy
	// AST-global storage to use for operations that are not supported by views
	// (i.e. fingerprint->metric lookups).
	storage *metric.TieredStorage
	// The materialized view which contains all timeseries data required for
	// executing a query.
	view metric.View
}

// interpolateSamples interpolates a value at a target time between two
// provided sample pairs.
func interpolateSamples(first, second *model.SamplePair, timestamp time.Time) *model.SamplePair {
	dv := second.Value - first.Value
	dt := second.Timestamp.Sub(first.Timestamp)

	dDt := dv / model.SampleValue(dt)
	offset := model.SampleValue(timestamp.Sub(first.Timestamp))

	return &model.SamplePair{
		Value:     first.Value + (offset * dDt),
		Timestamp: timestamp,
	}
}

// chooseClosestSample chooses the closest sample of a list of samples
// surrounding a given target time. If samples are found both before and after
// the target time, the sample value is interpolated between these. Otherwise,
// the single closest sample is returned verbatim.
func (v *viewAdapter) chooseClosestSample(samples model.Values, timestamp time.Time) *model.SamplePair {
	var closestBefore *model.SamplePair
	var closestAfter *model.SamplePair
	for _, candidate := range samples {
		delta := candidate.Timestamp.Sub(timestamp)
		// Samples before target time.
		if delta < 0 {
			// Ignore samples outside of staleness policy window.
			if -delta > v.stalenessPolicy.DeltaAllowance {
				continue
			}
			// Ignore samples that are farther away than what we've seen before.
			if closestBefore != nil && candidate.Timestamp.Before(closestBefore.Timestamp) {
				continue
			}
			sample := candidate
			closestBefore = &sample
		}

		// Samples after target time.
		if delta >= 0 {
			// Ignore samples outside of staleness policy window.
			if delta > v.stalenessPolicy.DeltaAllowance {
				continue
			}
			// Ignore samples that are farther away than samples we've seen before.
			if closestAfter != nil && candidate.Timestamp.After(closestAfter.Timestamp) {
				continue
			}
			sample := candidate
			closestAfter = &sample
		}
	}

	switch {
	case closestBefore != nil && closestAfter != nil:
		return interpolateSamples(closestBefore, closestAfter, timestamp)
	case closestBefore != nil:
		return closestBefore
	default:
		return closestAfter
	}
}

func (v *viewAdapter) GetValueAtTime(fingerprints model.Fingerprints, timestamp time.Time) (samples Vector, err error) {
	for _, fingerprint := range fingerprints {
		sampleCandidates := v.view.GetValueAtTime(fingerprint, timestamp)
		samplePair := v.chooseClosestSample(sampleCandidates, timestamp)
		m, err := v.storage.GetMetricForFingerprint(fingerprint)
		if err != nil {
			continue
		}
		if samplePair != nil {
			samples = append(samples, model.Sample{
				Metric:    m,
				Value:     samplePair.Value,
				Timestamp: timestamp,
			})
		}
	}
	return samples, err
}

func (v *viewAdapter) GetBoundaryValues(fingerprints model.Fingerprints, interval *model.Interval) (sampleSets []model.SampleSet, err error) {
	for _, fingerprint := range fingerprints {
		samplePairs := v.view.GetBoundaryValues(fingerprint, *interval)
		if len(samplePairs) == 0 {
			continue
		}

		// TODO: memoize/cache this.
		m, err := v.storage.GetMetricForFingerprint(fingerprint)
		if err != nil {
			continue
		}

		sampleSet := model.SampleSet{
			Metric: m,
			Values: samplePairs,
		}
		sampleSets = append(sampleSets, sampleSet)
	}
	return sampleSets, nil
}

func (v *viewAdapter) GetRangeValues(fingerprints model.Fingerprints, interval *model.Interval) (sampleSets []model.SampleSet, err error) {
	for _, fingerprint := range fingerprints {
		samplePairs := v.view.GetRangeValues(fingerprint, *interval)
		if len(samplePairs) == 0 {
			continue
		}

		// TODO: memoize/cache this.
		m, err := v.storage.GetMetricForFingerprint(fingerprint)
		if err != nil {
			continue
		}

		sampleSet := model.SampleSet{
			Metric: m,
			Values: samplePairs,
		}
		sampleSets = append(sampleSets, sampleSet)
	}
	return sampleSets, nil
}

func NewViewAdapter(view metric.View, storage *metric.TieredStorage) *viewAdapter {
	stalenessPolicy := StalenessPolicy{
		DeltaAllowance: time.Duration(*defaultStalenessDelta) * time.Second,
	}

	return &viewAdapter{
		stalenessPolicy: stalenessPolicy,
		storage:         storage,
		view:            view,
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/zyedidia/micro/cmd/micro/highlight"
)

func main() {
	files, _ := ioutil.ReadDir(".")

	hadErr := false
	for _, f := range files {
		if strings.HasSuffix(f.Name(), ".yaml") {
			input, _ := ioutil.ReadFile(f.Name())
			_, err := highlight.ParseDef(input)
			if err != nil {
				hadErr = true
				fmt.Printf("%s:\n", f.Name())
				fmt.Println(err)
				continue
			}
		}
	}
	if !hadErr {
		fmt.Println("No issues!")
	}
}
<commit_msg>make syntax_checker work with error messages<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/zyedidia/micro/cmd/micro/highlight"
)

func main() {
	files, _ := ioutil.ReadDir(".")

	hadErr := false
	for _, f := range files {
		if strings.HasSuffix(f.Name(), ".yaml") {
			input, _ := ioutil.ReadFile(f.Name())
			//fmt.Println("Checking file -> ", f.Name())
			file, err := highlight.ParseFile(input)
			if err != nil {
				hadErr = true
				fmt.Printf("Could not parse file -> %s:\n", f.Name())
				fmt.Println(err)
				continue
			}
			_, err1 := highlight.ParseDef(file, nil)
			if err1 != nil {
				hadErr = true
				fmt.Printf("Could not parse input file using highlight.ParseDef(%s):\n", f.Name())
				fmt.Println(err1)
				continue
			}
		}
	}
	if !hadErr {
		fmt.Println("No issues found!")
	}
}
<|endoftext|>