<commit_before>package http

import (
	"github.com/creamdog/aproxy/mappings"
	"github.com/creamdog/aproxy/cache"
	"io"
	"log"
	"net/http"
	"strings"
	"encoding/json"
	"bytes"
	"fmt"
	"time"
)

type HttpPipe struct {
	cache cache.CacheClient
}

type CachedResponse struct {
	Header map[string][]string
	StatusCode int
	Body string
	Expires int
	Key string
}

func New(cacheClient cache.CacheClient) *HttpPipe {
	return &HttpPipe{
		cache : cacheClient,
	}
}

func (pipe *HttpPipe) Pipe(mapping *mappings.RequestMapping, w http.ResponseWriter) {

	log.Printf("%v %v, %v, cacheKey: %s", mapping.Verb, mapping.Uri, mapping.Headers, mapping.CacheKey)

	w.Header().Set("X-AProxy-Version", "0.1")
	_, notransform := (*mapping.Data)["query"].(map[string]interface{})["_notransform"]
	_, nocache := (*mapping.Data)["query"].(map[string]interface{})["_nocache"]

	if nocache {
		mapping.CacheKey = ""
	}

	if len(mapping.CacheKey) > 0 {
		var cacheResponse CachedResponse
		if ok, err := pipe.cache.Get(mapping.CacheKey, &cacheResponse); ok {
			log.Printf("cache hit: %v", mapping.CacheKey)
			for key, value := range cacheResponse.Header {
				w.Header().Set(key, value[0])
			}

			w.Header().Set("X-Cache-Hit", "true")
			w.Header().Set("X-Cache-Key", cacheResponse.Key)
			w.Header().Set("X-Cache-Expiration-Seconds", fmt.Sprintf("%d", int64(cacheResponse.Expires) - time.Now().Unix()))

			w.WriteHeader(cacheResponse.StatusCode)
			fmt.Fprint(w, cacheResponse.Body)
			return
		} else if err != nil {
			log.Printf("%v", err)
		} else {
			log.Printf("cache miss '%s'", mapping.CacheKey)
		}
	}

	request, err := http.NewRequest(mapping.Verb, mapping.Uri, strings.NewReader(mapping.Body))

	if err != nil {
		http.Error(w, err.Error(), 503)
	}
	request.ContentLength = int64(len(mapping.Body))

	//log.Printf("request.ContentLength: %d, mapping.Body: %v", request.ContentLength, mapping.Body)

	for key, value := range mapping.Headers {
		request.Header[key] = []string{value}
	}

	request.Header["Transfer-Encoding"] = []string{""}

	client := &http.Client{}
	if response, err := client.Do(request); err != nil {
		http.Error(w, err.Error(), 500)
	} else {

		defer response.Body.Close()

		responseBody := ""
		responseBodyRead := false

		readResponseBody := func() ([]byte, error) {
			responseBodyRead = true
			readBuffer := make([]byte, 1024)
			buffer := make([]byte, 0)
			read := true
			//max := 1 * 1024 * 1024
			for read {
				/*
				if len(buffer) >= max {
					http.Error(w, fmt.Errorf("transform: maximum input site exceeded %d bytes", max).Error(), 500)
					return
				}
				*/
				read = false
				num, err := response.Body.Read(readBuffer)
				if err != nil {
					if err.Error() == "EOF" {
						buffer = append(buffer, readBuffer[:num]...)
						break
					}
					return nil, err
				}
				//log.Printf("read %d bytes", num)
				if num > 0 {
					buffer = append(buffer, readBuffer[:num]...)
					read = true
				}
			}
			return buffer, nil
		}

		if mapping.CompiledTransform != nil && !notransform {

			buffer, err := readResponseBody()
			if err != nil {
				http.Error(w, err.Error(), 500)
			}

			//log.Printf("buffer[%d]: %v", len(buffer), string(buffer))

			responseData := map[string]interface{}{}

			if mapping.Mapping.Target.Transform.Type == "json" {
				if err := json.Unmarshal(buffer, &responseData); err != nil {
					http.Error(w, err.Error() + " : " + string(buffer), 500)
					return
				}
			} else if mapping.Mapping.Target.Transform.Type == "regexp" {
				re := mapping.Mapping.Target.Transform.Regexp.FindStringSubmatch(string(buffer))
				names := mapping.Mapping.Target.Transform.Regexp.SubexpNames()
				if re != nil {
					for i, n := range re {
						if len(names[i]) > 0 {
							responseData[names[i]] = n
						}
					}
				}
			}

			//log.Printf("responseData: %v", responseData)

			data := map[string]interface{}{
				"data" : responseData,
			}
			for key, value := range *mapping.Data {
				data[key] = value
			}

			var renderBuffer bytes.Buffer
			if err := mapping.CompiledTransform.Execute(&renderBuffer, data); err != nil {
				http.Error(w, err.Error(), 500)
				return
			}
			responseBody = renderBuffer.String()
		}

		for key, value := range response.Header {
			w.Header().Set(key, value[0])
		}
		w.WriteHeader(response.StatusCode)

		if len(mapping.CacheKey) > 0 && !responseBodyRead {
			buffer, err := readResponseBody()
			if err != nil {
				http.Error(w, err.Error(), 500)
			}
			responseBody = string(buffer)
		}

		if responseBodyRead {

			if len(mapping.CacheKey) > 0 {
				cachedResponse := CachedResponse{
					Header : response.Header,
					StatusCode: response.StatusCode,
					Body: responseBody,
					Expires: int(time.Now().Unix()) + mapping.Mapping.Caching.Seconds,
					Key: mapping.CacheKey,
				}
				pipe.cache.Set(mapping.CacheKey, cachedResponse.Expires, cachedResponse)
			}

			fmt.Fprint(w, responseBody)
		} else {
			io.Copy(w, response.Body)
		}
	}
}
<commit_msg>override content-length on content transformation<commit_after>package http

import (
	"github.com/creamdog/aproxy/mappings"
	"github.com/creamdog/aproxy/cache"
	"io"
	"log"
	"net/http"
	"strings"
	"encoding/json"
	"bytes"
	"fmt"
	"time"
)

type HttpPipe struct {
	cache cache.CacheClient
}

type CachedResponse struct {
	Header map[string][]string
	StatusCode int
	Body string
	Expires int
	Key string
}

func New(cacheClient cache.CacheClient) *HttpPipe {
	return &HttpPipe{
		cache : cacheClient,
	}
}

func (pipe *HttpPipe) Pipe(mapping *mappings.RequestMapping, w http.ResponseWriter) {

	log.Printf("%v %v, %v, cacheKey: %s", mapping.Verb, mapping.Uri, mapping.Headers, mapping.CacheKey)

	w.Header().Set("X-AProxy-Version", "0.1")
	_, notransform := (*mapping.Data)["query"].(map[string]interface{})["_notransform"]
	_, nocache := (*mapping.Data)["query"].(map[string]interface{})["_nocache"]

	maxBodySize := 1 * 1024 * 1024

	if nocache {
		mapping.CacheKey = ""
	}

	if len(mapping.CacheKey) > 0 {
		var cacheResponse CachedResponse
		if ok, err := pipe.cache.Get(mapping.CacheKey, &cacheResponse); ok {
			log.Printf("cache hit: %v", mapping.CacheKey)
			for key, value := range cacheResponse.Header {
				if key == "Content-Length" {
					w.Header().Set("Content-Length", fmt.Sprintf("%d", len(cacheResponse.Body)))
				} else {
					w.Header().Set(key, value[0])
				}
			}

			w.Header().Set("X-Cache-Hit", "true")
			w.Header().Set("X-Cache-Key", cacheResponse.Key)
			w.Header().Set("X-Cache-Expiration-Seconds", fmt.Sprintf("%d", int64(cacheResponse.Expires) - time.Now().Unix()))

			w.WriteHeader(cacheResponse.StatusCode)
			fmt.Fprint(w, cacheResponse.Body)
			return
		} else if err != nil {
			log.Printf("%v", err)
		} else {
			log.Printf("cache miss '%s'", mapping.CacheKey)
		}
	}

	request, err := http.NewRequest(mapping.Verb, mapping.Uri, strings.NewReader(mapping.Body))

	if err != nil {
		http.Error(w, err.Error(), 503)
	}
	request.ContentLength = int64(len(mapping.Body))

	//log.Printf("request.ContentLength: %d, mapping.Body: %v", request.ContentLength, mapping.Body)

	for key, value := range mapping.Headers {
		request.Header[key] = []string{value}
	}

	request.Header["Transfer-Encoding"] = []string{""}

	client := &http.Client{}
	if response, err := client.Do(request); err != nil {
		http.Error(w, err.Error(), 500)
	} else {

		defer response.Body.Close()

		responseBody := ""
		responseBodyRead := false

		readResponseBody := func() ([]byte, error) {
			responseBodyRead = true
			readBuffer := make([]byte, 1024)
			buffer := make([]byte, 0)
			read := true
			//max := 1 * 1024 * 1024
			for read {

				if len(buffer) > maxBodySize {

				}

				/*
				if len(buffer) >= max {
					http.Error(w, fmt.Errorf("transform: maximum input site exceeded %d bytes", max).Error(), 500)
					return
				}
				*/
				read = false
				num, err := response.Body.Read(readBuffer)
				if err != nil {
					if err.Error() == "EOF" {
						buffer = append(buffer, readBuffer[:num]...)
						break
					}
					return nil, err
				}
				//log.Printf("read %d bytes", num)
				if num > 0 {
					buffer = append(buffer, readBuffer[:num]...)
					read = true
				}
			}
			return buffer, nil
		}

		if mapping.CompiledTransform != nil && !notransform && response.StatusCode == 200 {

			buffer, err := readResponseBody()
			if err != nil {
				http.Error(w, err.Error(), 500)
			}

			//log.Printf("buffer[%d]: %v", len(buffer), string(buffer))

			responseData := map[string]interface{}{}

			if mapping.Mapping.Target.Transform.Type == "json" {
				if err := json.Unmarshal(buffer, &responseData); err != nil {
					http.Error(w, err.Error() + " : " + string(buffer), 500)
					return
				}
			} else if mapping.Mapping.Target.Transform.Type == "regexp" {
				re := mapping.Mapping.Target.Transform.Regexp.FindStringSubmatch(string(buffer))
				names := mapping.Mapping.Target.Transform.Regexp.SubexpNames()
				if re != nil {
					for i, n := range re {
						if len(names[i]) > 0 {
							responseData[names[i]] = n
						}
					}
				}
			}

			//log.Printf("responseData: %v", responseData)

			data := map[string]interface{}{
				"data" : responseData,
			}
			for key, value := range *mapping.Data {
				data[key] = value
			}

			var renderBuffer bytes.Buffer
			if err := mapping.CompiledTransform.Execute(&renderBuffer, data); err != nil {
				http.Error(w, err.Error(), 500)
				return
			}
			responseBody = renderBuffer.String()
		}

		if len(mapping.CacheKey) > 0 && !responseBodyRead {
			buffer, err := readResponseBody()
			if err != nil {
				http.Error(w, err.Error(), 500)
			}
			responseBody = string(buffer)
		}

		for key, value := range response.Header {
			if responseBodyRead && key == "Content-Length" {
				w.Header().Set("Content-Length", fmt.Sprintf("%d", len(responseBody)))
			} else {
				w.Header().Set(key, value[0])
			}
		}

		w.WriteHeader(response.StatusCode)

		if responseBodyRead {

			if len(mapping.CacheKey) > 0 {
				cachedResponse := CachedResponse{
					Header : response.Header,
					StatusCode: response.StatusCode,
					Body: responseBody,
					Expires: int(time.Now().Unix()) + mapping.Mapping.Caching.Seconds,
					Key: mapping.CacheKey,
				}
				pipe.cache.Set(mapping.CacheKey, cachedResponse.Expires, cachedResponse)
			}

			fmt.Fprint(w, responseBody)
		} else {
			io.Copy(w, response.Body)
		}
	}
}
<|endoftext|>


<commit_before>// Brief: Bible Passage handling
// Primary responsibility: Parsing of Bible Passages from HTML

package app

import (
	"fmt"
	"log"
	"strings"

	"golang.org/x/net/html"

	"github.com/julwrites/BotMultiplexer/pkg/def"
	"github.com/julwrites/ScriptureBot/pkg/utils"
)

func GetReference(doc *html.Node) string {
	refNode, err := utils.FindByClass(doc, "bcv")
	if err != nil {
		log.Printf("Error parsing for reference: %v", err)
		return ""
	}

	return refNode.FirstChild.Data
}

func ParseNodesForPassage(node *html.Node) string {
	var text string
	var parts []string

	for child := node.FirstChild; child != nil; child = child.NextSibling {
		parts = append(parts, text)

		switch tag := child.Data; tag {
		case "span":
			childText := strings.Trim(ParseNodesForPassage(child), " ")
			if len(childText) > 0 {
				parts = append(parts, childText)
			}
		case "sup":
			isFootnote := func(node *html.Node) bool {
				for _, attr := range node.Attr {
					if attr.Key == "class" && attr.Val == "footnote" {
						return true
					}
				}
				return false
			}
			if isFootnote(child) {
				break
			}
			childText := strings.Trim(ParseNodesForPassage(child), " ")
			if len(childText) > 0 {
				parts = append(parts, fmt.Sprintf("^%s^", childText))
			}
			break
		case "p":
			parts = append(parts, ParseNodesForPassage(child))
			break
		case "br":
			parts = append(parts, " ")
			break
		default:
			parts = append(parts, child.Data)
		}
	}

	text = strings.Join(parts, "")

	if node.Data == "h1" || node.Data == "h2" || node.Data == "h3" || node.Data == "h4" {
		text = fmt.Sprintf("*%s*", text)
	}
	return text
}

func GetPassage(doc *html.Node, env *def.SessionData) string {
	passageNode, startErr := utils.FindByClass(doc, fmt.Sprintf("version-%s result-text-style-normal text-html", utils.DeserializeUserConfig(env.User.Config).Version))
	if startErr != nil {
		log.Printf("Error parsing for passage: %v", startErr)
		return ""
	}

	filtNodes := utils.FilterChildren(passageNode, func(child *html.Node) bool {
		switch tag := child.Data; tag {
		case "h1":
			fallthrough
		case "h2":
			fallthrough
		case "h3":
			fallthrough
		case "h4":
			fallthrough
		case "p":
			return true
		case "div":
			for _, attr := range child.Attr {
				if attr.Key == "class" && strings.Contains(attr.Val, "poetry") {
					return true
				}
			}
		}
		return false
	})

	textBlocks := utils.MapNodeList(filtNodes, ParseNodesForPassage)

	var passage strings.Builder

	for _, block := range textBlocks {
		passage.WriteString(block)
		passage.WriteString("\n")
	}

	return passage.String()
}

func GetBiblePassage(env def.SessionData) def.SessionData {
	if len(env.Msg.Message) > 0 {

		doc := utils.QueryBiblePassage(env.Msg.Message, utils.DeserializeUserConfig(env.User.Config).Version)

		ref := GetReference(doc)
		log.Printf("Reference retrieved: %s", ref)

		if len(ref) > 0 {
			log.Printf("Getting passage")
			env.Res.Message = GetPassage(doc, &env)
			log.Printf("Passage retrieved %s", env.Res.Message)
		}
	}

	return env
}
<commit_msg>Changing back to newline<commit_after>// Brief: Bible Passage handling
// Primary responsibility: Parsing of Bible Passages from HTML

package app

import (
	"fmt"
	"log"
	"strings"

	"golang.org/x/net/html"

	"github.com/julwrites/BotMultiplexer/pkg/def"
	"github.com/julwrites/ScriptureBot/pkg/utils"
)

func GetReference(doc *html.Node) string {
	refNode, err := utils.FindByClass(doc, "bcv")
	if err != nil {
		log.Printf("Error parsing for reference: %v", err)
		return ""
	}

	return refNode.FirstChild.Data
}

func ParseNodesForPassage(node *html.Node) string {
	var text string
	var parts []string

	for child := node.FirstChild; child != nil; child = child.NextSibling {
		parts = append(parts, text)

		switch tag := child.Data; tag {
		case "span":
			childText := strings.Trim(ParseNodesForPassage(child), " ")
			if len(childText) > 0 {
				parts = append(parts, childText)
			}
		case "sup":
			isFootnote := func(node *html.Node) bool {
				for _, attr := range node.Attr {
					if attr.Key == "class" && attr.Val == "footnote" {
						return true
					}
				}
				return false
			}
			if isFootnote(child) {
				break
			}
			childText := strings.Trim(ParseNodesForPassage(child), " ")
			if len(childText) > 0 {
				parts = append(parts, fmt.Sprintf("^%s^", childText))
			}
			break
		case "p":
			parts = append(parts, ParseNodesForPassage(child))
			break
		case "br":
			parts = append(parts, "\n")
			break
		default:
			parts = append(parts, child.Data)
		}
	}

	text = strings.Join(parts, "")

	if node.Data == "h1" || node.Data == "h2" || node.Data == "h3" || node.Data == "h4" {
		text = fmt.Sprintf("*%s*", text)
	}
	return text
}

func GetPassage(doc *html.Node, env *def.SessionData) string {
	passageNode, startErr := utils.FindByClass(doc, fmt.Sprintf("version-%s result-text-style-normal text-html", utils.DeserializeUserConfig(env.User.Config).Version))
	if startErr != nil {
		log.Printf("Error parsing for passage: %v", startErr)
		return ""
	}

	filtNodes := utils.FilterChildren(passageNode, func(child *html.Node) bool {
		switch tag := child.Data; tag {
		case "h1":
			fallthrough
		case "h2":
			fallthrough
		case "h3":
			fallthrough
		case "h4":
			fallthrough
		case "p":
			return true
		case "div":
			for _, attr := range child.Attr {
				if attr.Key == "class" && strings.Contains(attr.Val, "poetry") {
					return true
				}
			}
		}
		return false
	})

	textBlocks := utils.MapNodeList(filtNodes, ParseNodesForPassage)

	var passage strings.Builder

	for _, block := range textBlocks {
		passage.WriteString(block)
		passage.WriteString("\n")
	}

	return passage.String()
}

func GetBiblePassage(env def.SessionData) def.SessionData {
	if len(env.Msg.Message) > 0 {

		doc := utils.QueryBiblePassage(env.Msg.Message, utils.DeserializeUserConfig(env.User.Config).Version)

		ref := GetReference(doc)
		log.Printf("Reference retrieved: %s", ref)

		if len(ref) > 0 {
			log.Printf("Getting passage")
			env.Res.Message = GetPassage(doc, &env)
			log.Printf("Passage retrieved %s", env.Res.Message)
		}
	}

	return env
}
<|endoftext|>


<commit_before>package ec2

//-----------------------------------------------------------------------------
// Package factored import statement:
//-----------------------------------------------------------------------------

import (

	// Stdlib:
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"time"

	// Community:
	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/elb"
)

//-----------------------------------------------------------------------------
// func: Run
//-----------------------------------------------------------------------------

// Run uses EC2 API to launch a new instance.
func (d *Data) Run() {

	// Set current command:
	d.command = "run"

	// Read udata from stdin:
	udata, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Connect and authenticate to the API endpoints:
	d.ec2 = ec2.New(session.New(&aws.Config{Region: aws.String(d.Region)}))
	d.elb = elb.New(session.New(&aws.Config{Region: aws.String(d.Region)}))

	// Run the EC2 instance:
	if err := d.runInstance(udata); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Modify instance attributes:
	if err := d.modifyInstanceAttribute(); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Setup an elastic IP:
	if d.PublicIP == "elastic" {
		if err := d.setupElasticIP(); err != nil {
			log.WithField("cmd", "ec2:"+d.command).Fatal(err)
		}
	}

	// Register with ELB:
	if d.ELBName != "" {
		if err := d.registerWithELB(); err != nil {
			log.WithField("cmd", "ec2:"+d.command).Fatal(err)
		}
	}

	// Output IP addresses to stdout:
	if err := d.stdoutIPs(); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Warning(err)
	}
}

//-----------------------------------------------------------------------------
// func: runInstance
//-----------------------------------------------------------------------------

func (d *Data) runInstance(udata []byte) error {

	// Variables:
	var resp *ec2.Reservation
	var err error

	// Forge the instance request:
	params := &ec2.RunInstancesInput{
		ImageId:           aws.String(d.AmiID),
		MinCount:          aws.Int64(1),
		MaxCount:          aws.Int64(1),
		KeyName:           aws.String(d.KeyPair),
		InstanceType:      aws.String(d.InstanceType),
		NetworkInterfaces: d.forgeNetworkInterfaces(),
		Placement: &ec2.Placement{
			AvailabilityZone: aws.String(d.Region + d.Zone),
		},
		UserData: aws.String(base64.StdEncoding.EncodeToString([]byte(udata))),
		IamInstanceProfile: &ec2.IamInstanceProfileSpecification{
			Name: aws.String(d.IAMRole),
		},
	}

	// Send the instance request:
	for i := 0; i < 5; i++ {
		resp, err = d.ec2.RunInstances(params)
		if err != nil {
			ec2err, ok := err.(awserr.Error)
			if ok && strings.Contains(ec2err.Code(), "InvalidParameterValue") {
				time.Sleep(2 * time.Second)
				continue
			}
			return err
		}
		break
	}

	// Store the instance ID:
	d.InstanceID = *resp.Instances[0].InstanceId
	log.WithFields(log.Fields{"cmd": "ec2:" + d.command, "id": d.InstanceID}).
		Info("New " + d.InstanceType + " EC2 instance requested")

	// Store the interface ID:
	d.InterfaceID = *resp.Instances[0].
		NetworkInterfaces[0].NetworkInterfaceId

	// Tag the instance:
	if err := d.tag(d.InstanceID, "Name", d.TagName); err != nil {
		return err
	}

	// Pretty-print to stderr:
	log.WithFields(log.Fields{"cmd": "ec2:" + d.command, "id": d.TagName}).
		Info("New EC2 instance tagged")

	return nil
}

//-----------------------------------------------------------------------------
// func: forgeNetworkInterfaces
//-----------------------------------------------------------------------------

func (d *Data) forgeNetworkInterfaces() []*ec2.
	InstanceNetworkInterfaceSpecification {

	var networkInterfaces []*ec2.InstanceNetworkInterfaceSpecification
	var securityGroupIds []*string

	// Append to security group array:
	for _, grp := range strings.Split(d.SecGrpIDs, ",") {
		securityGroupIds = append(securityGroupIds, aws.String(grp))
	}

	// Forge the interface data type:
	iface := ec2.InstanceNetworkInterfaceSpecification{
		DeleteOnTermination: aws.Bool(true),
		DeviceIndex:         aws.Int64(int64(0)),
		Groups:              securityGroupIds,
		SubnetId:            aws.String(d.SubnetID),
	}

	// Private IP address:
	if d.PrivateIP != "" {
		iface.PrivateIpAddress = aws.String(d.PrivateIP)
	}

	// Public IP address:
	if d.PublicIP == "true" {
		iface.AssociatePublicIpAddress = aws.Bool(true)
	}

	// Append to the interfaces array:
	networkInterfaces = append(networkInterfaces, &iface)

	return networkInterfaces
}

//-----------------------------------------------------------------------------
// func: modifyInstanceAttribute
//-----------------------------------------------------------------------------

func (d *Data) modifyInstanceAttribute() error {

	// Variable transformation:
	SrcDstCheck, err := strconv.ParseBool(d.SrcDstCheck)
	if err != nil {
		return err
	}

	// Forge the attribute modification request:
	params := &ec2.ModifyInstanceAttributeInput{
		InstanceId: aws.String(d.InstanceID),
		SourceDestCheck: &ec2.AttributeBooleanValue{
			Value: aws.Bool(SrcDstCheck),
		},
	}

	// Send the attribute modification request:
	_, err = d.ec2.ModifyInstanceAttribute(params)
	if err != nil {
		return err
	}

	return nil
}

//-----------------------------------------------------------------------------
// func: setupElasticIP
//-----------------------------------------------------------------------------

func (d *Data) setupElasticIP() error {

	// Allocate an elastic IP address:
	if err := d.allocateElasticIP(); err != nil {
		return err
	}

	// Associate the elastic IP:
	if err := d.associateElasticIP(); err != nil {
		return err
	}

	return nil
}

//-----------------------------------------------------------------------------
// func: associateElasticIP
//-----------------------------------------------------------------------------

func (d *Data) associateElasticIP() error {

	// Wait until instance is running:
	if err := d.ec2.WaitUntilInstanceRunning(&ec2.DescribeInstancesInput{
		InstanceIds: []*string{aws.String(d.InstanceID)}}); err != nil {
		return err
	}

	// Forge the association request:
	params := &ec2.AssociateAddressInput{
		AllocationId:       aws.String(d.AllocationID),
		AllowReassociation: aws.Bool(true),
		NetworkInterfaceId: aws.String(d.InterfaceID),
	}

	// Send the association request:
	resp, err := d.ec2.AssociateAddress(params)
	if err != nil {
		return err
	}

	// Log to stderr:
	log.WithFields(log.Fields{
		"cmd": "ec2:" + d.command, "id": *resp.AssociationId}).
		Info("New elastic IP association")

	return nil
}

//-----------------------------------------------------------------------------
// func: registerWithELB
//-----------------------------------------------------------------------------

func (d *Data) registerWithELB() error {

	// Forge the register request:
	params := &elb.RegisterInstancesWithLoadBalancerInput{
		Instances: []*elb.Instance{
			{
				InstanceId: aws.String(d.InstanceID),
			},
		},
		LoadBalancerName: aws.String(d.ELBName),
	}

	// Send the register request:
	if _, err := d.elb.RegisterInstancesWithLoadBalancer(params); err != nil {
		return err
	}

	// Log this action:
	log.WithFields(log.Fields{"cmd": "ec2:" + d.command, "id": d.ELBName}).
		Info("Instance registered with ELB")

	return nil
}

//-----------------------------------------------------------------------------
// func: stdoutIPs
//-----------------------------------------------------------------------------

func (d *Data) stdoutIPs() error {

	// Map to store the output:
	m := make(map[string]string)

	// Forge the describe request:
	params := &ec2.DescribeNetworkInterfacesInput{
		NetworkInterfaceIds: []*string{
			aws.String(d.InterfaceID),
		},
	}

	// Retry loop:
	for i := 0; i < 5; i++ {

		// Send the describe request:
		resp, err := d.ec2.DescribeNetworkInterfaces(params)
		if err != nil {
			return err
		}

		// Extract data from response:
		if len(resp.NetworkInterfaces) > 0 && len(resp.NetworkInterfaces[0].PrivateIpAddresses) > 0 {

			// Internal IP address:
			if resp.NetworkInterfaces[0].PrivateIpAddresses[0].PrivateIpAddress != nil {
				m["internal"] = *resp.NetworkInterfaces[0].PrivateIpAddresses[0].PrivateIpAddress
			}

			// External IP address:
			if resp.NetworkInterfaces[0].PrivateIpAddresses[0].Association != nil {
				if resp.NetworkInterfaces[0].PrivateIpAddresses[0].Association.PublicIp != nil {
					m["external"] = *resp.NetworkInterfaces[0].PrivateIpAddresses[0].Association.PublicIp
				}
			}
		}

		// Break if IP:
		if m["internal"] != "" {
			break
		}

		// Sleep and try again:
		time.Sleep(2 * time.Second)
	}

	// JSON encode:
	jsn, err := json.Marshal(m)
	if err != nil {
		return err
	}

	// Print and return:
	fmt.Println(string(jsn))
	return nil
}
<commit_msg>Same thing different code<commit_after>package ec2

//-----------------------------------------------------------------------------
// Package factored import statement:
//-----------------------------------------------------------------------------

import (

	// Stdlib:
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"time"

	// Community:
	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/elb"
)

//-----------------------------------------------------------------------------
// func: Run
//-----------------------------------------------------------------------------

// Run uses EC2 API to launch a new instance.
func (d *Data) Run() {

	// Set current command:
	d.command = "run"

	// Read udata from stdin:
	udata, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Connect and authenticate to the API endpoints:
	d.ec2 = ec2.New(session.New(&aws.Config{Region: aws.String(d.Region)}))
	d.elb = elb.New(session.New(&aws.Config{Region: aws.String(d.Region)}))

	// Run the EC2 instance:
	if err := d.runInstance(udata); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Modify instance attributes:
	if err := d.modifyInstanceAttribute(); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Fatal(err)
	}

	// Setup an elastic IP:
	if d.PublicIP == "elastic" {
		if err := d.setupElasticIP(); err != nil {
			log.WithField("cmd", "ec2:"+d.command).Fatal(err)
		}
	}

	// Register with ELB:
	if d.ELBName != "" {
		if err := d.registerWithELB(); err != nil {
			log.WithField("cmd", "ec2:"+d.command).Fatal(err)
		}
	}

	// Output IP addresses to stdout:
	if err := d.stdoutIPs(); err != nil {
		log.WithField("cmd", "ec2:"+d.command).Warning(err)
	}
}

//-----------------------------------------------------------------------------
// func: runInstance
//-----------------------------------------------------------------------------

func (d *Data) runInstance(udata []byte) error {

	// Variables:
	var resp *ec2.Reservation
	var err error

	// Forge the instance request:
	params := &ec2.RunInstancesInput{
		ImageId:           aws.String(d.AmiID),
		MinCount:          aws.Int64(1),
		MaxCount:          aws.Int64(1),
		KeyName:           aws.String(d.KeyPair),
		InstanceType:      aws.String(d.InstanceType),
		NetworkInterfaces: d.forgeNetworkInterfaces(),
		Placement: &ec2.Placement{
			AvailabilityZone: aws.String(d.Region + d.Zone),
		},
		UserData: aws.String(base64.StdEncoding.EncodeToString([]byte(udata))),
		IamInstanceProfile: &ec2.IamInstanceProfileSpecification{
			Name: aws.String(d.IAMRole),
		},
	}

	// Send the instance request:
	for i := 0; i < 5; i++ {
		resp, err = d.ec2.RunInstances(params)
		if err != nil {
			ec2err, ok := err.(awserr.Error)
			if ok && strings.Contains(ec2err.Code(), "InvalidParameterValue") {
				time.Sleep(2 * time.Second)
				continue
			}
			return err
		}
		break
	}

	// Store the instance ID:
	d.InstanceID = *resp.Instances[0].InstanceId
	log.WithFields(log.Fields{"cmd": "ec2:" + d.command, "id": d.InstanceID}).
		Info("New " + d.InstanceType + " EC2 instance requested")

	// Store the interface ID:
	d.InterfaceID = *resp.Instances[0].
		NetworkInterfaces[0].NetworkInterfaceId

	// Tag the instance:
	if err := d.tag(d.InstanceID, "Name", d.TagName); err != nil {
		return err
	}

	// Pretty-print to stderr:
	log.WithFields(log.Fields{"cmd": "ec2:" + d.command, "id": d.TagName}).
		Info("New EC2 instance tagged")

	return nil
}

//-----------------------------------------------------------------------------
// func: forgeNetworkInterfaces
//-----------------------------------------------------------------------------

func (d *Data) forgeNetworkInterfaces() []*ec2.
	InstanceNetworkInterfaceSpecification {

	var networkInterfaces []*ec2.InstanceNetworkInterfaceSpecification
	var securityGroupIds []*string

	// Append to security group array:
	for _, grp := range strings.Split(d.SecGrpIDs, ",") {
		securityGroupIds = append(securityGroupIds, aws.String(grp))
	}

	// Forge the interface data type:
	iface := ec2.InstanceNetworkInterfaceSpecification{
		DeleteOnTermination: aws.Bool(true),
		DeviceIndex:         aws.Int64(int64(0)),
		Groups:              securityGroupIds,
		SubnetId:            aws.String(d.SubnetID),
	}

	// Private IP address:
	if d.PrivateIP != "" {
		iface.PrivateIpAddress = aws.String(d.PrivateIP)
	}

	// Public IP address:
	if d.PublicIP == "true" {
		iface.AssociatePublicIpAddress = aws.Bool(true)
	}

	// Append to the interfaces array:
	networkInterfaces = append(networkInterfaces, &iface)

	return networkInterfaces
}

//-----------------------------------------------------------------------------
// func: modifyInstanceAttribute
//-----------------------------------------------------------------------------

func (d *Data) modifyInstanceAttribute() error {

	// Variable transformation:
	SrcDstCheck, err := strconv.ParseBool(d.SrcDstCheck)
	if err != nil {
		return err
	}

	// Forge the attribute modification request:
	params := &ec2.ModifyInstanceAttributeInput{
		InstanceId: aws.String(d.InstanceID),
		SourceDestCheck: &ec2.AttributeBooleanValue{
			Value: aws.Bool(SrcDstCheck),
		},
	}

	// Send the attribute modification request:
	_, err = d.ec2.ModifyInstanceAttribute(params)
	if err != nil {
		return err
	}

	return nil
}

//-----------------------------------------------------------------------------
// func: setupElasticIP
//-----------------------------------------------------------------------------

func (d *Data) setupElasticIP() error {

	// Allocate an elastic IP address:
	if err := d.allocateElasticIP(); err != nil {
		return err
	}

	// Associate the elastic IP:
	if err := d.associateElasticIP(); err != nil {
		return err
	}

	return nil
}

//-----------------------------------------------------------------------------
// func: associateElasticIP
//-----------------------------------------------------------------------------

func (d *Data) associateElasticIP() error {

	// Wait until instance is running:
	if err := d.ec2.WaitUntilInstanceRunning(&ec2.DescribeInstancesInput{
		InstanceIds: []*string{aws.String(d.InstanceID)}}); err != nil {
		return err
	}

	// Forge the association request:
	params := &ec2.AssociateAddressInput{
		AllocationId:       aws.String(d.AllocationID),
		AllowReassociation: aws.Bool(true),
		NetworkInterfaceId: aws.String(d.InterfaceID),
	}

	// Send the association request:
	resp, err := d.ec2.AssociateAddress(params)
	if err != nil {
		return err
	}

	// Log to stderr:
	log.WithFields(log.Fields{
		"cmd": "ec2:" + d.command, "id": *resp.AssociationId}).
		Info("New elastic IP association")

	return nil
}

//-----------------------------------------------------------------------------
// func: registerWithELB
//-----------------------------------------------------------------------------

func (d *Data) registerWithELB() error {

	// Forge the register request:
	params := &elb.RegisterInstancesWithLoadBalancerInput{
		Instances: []*elb.Instance{
			{
				InstanceId: aws.String(d.InstanceID),
			},
		},
		LoadBalancerName: aws.String(d.ELBName),
	}

	// Send the register request:
	if _, err := d.elb.RegisterInstancesWithLoadBalancer(params); err != nil {
		return err
	}

	// Log this action:
	log.WithFields(log.Fields{"cmd": "ec2:" + d.command, "id": d.ELBName}).
		Info("Instance registered with ELB")

	return nil
}

//-----------------------------------------------------------------------------
// func: stdoutIPs
//-----------------------------------------------------------------------------

func (d *Data) stdoutIPs() error {

	// Map to store the output:
	m := make(map[string]string)

	// Forge the describe request:
	params := &ec2.DescribeNetworkInterfacesInput{
		NetworkInterfaceIds: []*string{
			aws.String(d.InterfaceID),
		},
	}

	// Retry loop:
	for i := 0; i < 5; i++ {

		// Send the describe request:
		resp, err := d.ec2.DescribeNetworkInterfaces(params)
		if err != nil {
			return err
		}

		// Extract data from response:
		if len(resp.NetworkInterfaces) > 0 && len(resp.NetworkInterfaces[0].PrivateIpAddresses) > 0 {

			// Internal IP address:
			if resp.NetworkInterfaces[0].PrivateIpAddresses[0].PrivateIpAddress != nil {
				m["internal"] = *resp.NetworkInterfaces[0].PrivateIpAddresses[0].PrivateIpAddress
			}

			// External IP address:
			if resp.NetworkInterfaces[0].PrivateIpAddresses[0].Association != nil {
				if resp.NetworkInterfaces[0].PrivateIpAddresses[0].Association.PublicIp != nil {
					m["external"] = *resp.NetworkInterfaces[0].PrivateIpAddresses[0].Association.PublicIp
				}
			}
		}

		// Sleep and try again:
		if m["internal"] == "" {
			time.Sleep(2 * time.Second)
			continue
		}

		break
	}

	// JSON encode:
	jsn, err := json.Marshal(m)
	if err != nil {
		return err
	}

	// Print and return:
	fmt.Println(string(jsn))
	return nil
}
<|endoftext|>


<commit_before>// Package edit implements the line editor for Elvish.
//
// The line editor is based on the cli package, which implements a general,
// Elvish-agnostic line editor, and multiple "addon" packages. This package
// glues them together and provides Elvish bindings for them.
package edit

import (
	"fmt"
	"sync"

	"src.elv.sh/pkg/cli"
	"src.elv.sh/pkg/eval"
	"src.elv.sh/pkg/eval/vals"
	"src.elv.sh/pkg/eval/vars"
	"src.elv.sh/pkg/parse"
	"src.elv.sh/pkg/store"
)

// Editor is the interface line editor for Elvish.
type Editor struct {
	app cli.App
	ns  *eval.Ns

	excMutex sync.RWMutex
	excList  vals.List
}

// An interface that wraps notifyf and notifyError. It is only implemented by
// the *Editor type; functions may take a notifier instead of *Editor argument
// to make it clear that they do not depend on other parts of *Editor.
type notifier interface {
	notifyf(format string, args ...interface{})
	notifyError(ctx string, e error)
}

// NewEditor creates a new editor. The tty is used for input and output. The
// Evaler is used for syntax highlighting and completion, and the Store is
// used for saving and retrieving command and directory history.
func NewEditor(tty cli.TTY, ev *eval.Evaler, st store.Store) *Editor {
	// Declare the Editor with a nil App first; some initialization functions
	// require a notifier as an argument, but does not use it immediately.
	ed := &Editor{excList: vals.EmptyList}
	nb := eval.NsBuilder{}
	appSpec := cli.AppSpec{TTY: tty}

	hs, err := newHistStore(st)
	if err != nil {
		// TODO(xiaq): Report the error.
	}

	initHighlighter(&appSpec, ev)
	initMaxHeight(&appSpec, nb)
	initReadlineHooks(&appSpec, ev, nb)
	initAddCmdFilters(&appSpec, ev, nb, hs)
	initInsertAPI(&appSpec, ed, ev, nb)
	initPrompts(&appSpec, ed, ev, nb)
	ed.app = cli.NewApp(appSpec)

	initExceptionsAPI(ed, nb)
	initVarsAPI(ed, nb)
	initCommandAPI(ed, ev, nb)
	initListings(ed, ev, st, hs, nb)
	initNavigation(ed, ev, nb)
	initCompletion(ed, ev, nb)
	initHistWalk(ed, ev, hs, nb)
	initInstant(ed, ev, nb)
	initMinibuf(ed, ev, nb)

	initBufferBuiltins(ed.app, nb)
	initTTYBuiltins(ed.app, tty, nb)
	initMiscBuiltins(ed.app, nb)
	initStateAPI(ed.app, nb)
	initStoreAPI(ed.app, nb, hs)

	ed.ns = nb.Ns()
	evalDefaultBinding(ev, ed.ns)

	return ed
}

//elvdoc:var exceptions
//
// A list of exceptions thrown from callbacks such as prompts. Useful for
// examining tracebacks and other metadata.

func initExceptionsAPI(ed *Editor, nb eval.NsBuilder) {
	nb.Add("exceptions", vars.FromPtrWithMutex(&ed.excList, &ed.excMutex))
}

func evalDefaultBinding(ev *eval.Evaler, ns *eval.Ns) {
	src := parse.Source{Name: "[default bindings]", Code: defaultBindingsElv}
	err := ev.Eval(src, eval.EvalCfg{Global: ns})
	if err != nil {
		panic(err)
	}
}

// ReadCode reads input from the user.
func (ed *Editor) ReadCode() (string, error) {
	return ed.app.ReadCode()
}

// Ns returns a namespace for manipulating the editor from Elvish code.
//
// See https://elv.sh/ref/edit.html for the Elvish API.
func (ed *Editor) Ns() *eval.Ns {
	return ed.ns
}

func (ed *Editor) notifyf(format string, args ...interface{}) {
	ed.app.Notify(fmt.Sprintf(format, args...))
}

func (ed *Editor) notifyError(ctx string, e error) {
	if exc, ok := e.(eval.Exception); ok {
		ed.excMutex.Lock()
		defer ed.excMutex.Unlock()
		ed.excList = ed.excList.Cons(exc)
		ed.notifyf("[%v error] %v\n"+
			`see stack trace with "show $edit:exceptions[%d]"`,
			ctx, e, ed.excList.Len()-1)
	} else {
		ed.notifyf("[%v error] %v", ctx, e)
	}
}
<commit_msg>pkg/edit: Improve godoc.<commit_after>// Package edit implements the line editor for Elvish.
//
// The line editor is based on the cli package, which implements a general,
// Elvish-agnostic line editor, and multiple "addon" packages. This package
// glues them together and provides Elvish bindings for them.
package edit

import (
	"fmt"
	"sync"

	"src.elv.sh/pkg/cli"
	"src.elv.sh/pkg/eval"
	"src.elv.sh/pkg/eval/vals"
	"src.elv.sh/pkg/eval/vars"
	"src.elv.sh/pkg/parse"
	"src.elv.sh/pkg/store"
)

// Editor is the interactive line editor for Elvish.
type Editor struct {
	app cli.App
	ns  *eval.Ns

	excMutex sync.RWMutex
	excList  vals.List
}

// An interface that wraps notifyf and notifyError. It is only implemented by
// the *Editor type; functions may take a notifier instead of *Editor argument
// to make it clear that they do not depend on other parts of *Editor.
type notifier interface {
	notifyf(format string, args ...interface{})
	notifyError(ctx string, e error)
}

// NewEditor creates a new editor. The TTY is used for input and output. The
// Evaler is used for syntax highlighting, completion, and calling callbacks.
// The Store is used for saving and retrieving command and directory history.
func NewEditor(tty cli.TTY, ev *eval.Evaler, st store.Store) *Editor {
	// Declare the Editor with a nil App first; some initialization functions
	// require a notifier as an argument, but does not use it immediately.
	ed := &Editor{excList: vals.EmptyList}
	nb := eval.NsBuilder{}
	appSpec := cli.AppSpec{TTY: tty}

	hs, err := newHistStore(st)
	if err != nil {
		// TODO(xiaq): Report the error.
	}

	initHighlighter(&appSpec, ev)
	initMaxHeight(&appSpec, nb)
	initReadlineHooks(&appSpec, ev, nb)
	initAddCmdFilters(&appSpec, ev, nb, hs)
	initInsertAPI(&appSpec, ed, ev, nb)
	initPrompts(&appSpec, ed, ev, nb)
	ed.app = cli.NewApp(appSpec)

	initExceptionsAPI(ed, nb)
	initVarsAPI(ed, nb)
	initCommandAPI(ed, ev, nb)
	initListings(ed, ev, st, hs, nb)
	initNavigation(ed, ev, nb)
	initCompletion(ed, ev, nb)
	initHistWalk(ed, ev, hs, nb)
	initInstant(ed, ev, nb)
	initMinibuf(ed, ev, nb)

	initBufferBuiltins(ed.app, nb)
	initTTYBuiltins(ed.app, tty, nb)
	initMiscBuiltins(ed.app, nb)
	initStateAPI(ed.app, nb)
	initStoreAPI(ed.app, nb, hs)

	ed.ns = nb.Ns()
	evalDefaultBinding(ev, ed.ns)

	return ed
}

//elvdoc:var exceptions
//
// A list of exceptions thrown from callbacks such as prompts. Useful for
// examining tracebacks and other metadata.

func initExceptionsAPI(ed *Editor, nb eval.NsBuilder) {
	nb.Add("exceptions", vars.FromPtrWithMutex(&ed.excList, &ed.excMutex))
}

func evalDefaultBinding(ev *eval.Evaler, ns *eval.Ns) {
	src := parse.Source{Name: "[default bindings]", Code: defaultBindingsElv}
	err := ev.Eval(src, eval.EvalCfg{Global: ns})
	if err != nil {
		panic(err)
	}
}

// ReadCode reads input from the user.
func (ed *Editor) ReadCode() (string, error) {
	return ed.app.ReadCode()
}

// Ns returns a namespace for manipulating the editor from Elvish code.
//
// See https://elv.sh/ref/edit.html for the Elvish API.
func (ed *Editor) Ns() *eval.Ns {
	return ed.ns
}

func (ed *Editor) notifyf(format string, args ...interface{}) {
	ed.app.Notify(fmt.Sprintf(format, args...))
}

func (ed *Editor) notifyError(ctx string, e error) {
	if exc, ok := e.(eval.Exception); ok {
		ed.excMutex.Lock()
		defer ed.excMutex.Unlock()
		ed.excList = ed.excList.Cons(exc)
		ed.notifyf("[%v error] %v\n"+
			`see stack trace with "show $edit:exceptions[%d]"`,
			ctx, e, ed.excList.Len()-1)
	} else {
		ed.notifyf("[%v error] %v", ctx, e)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    https://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 flags defines global controller flags
package flags

import (
	"flag"
)

var (
	F = struct {
		APIServerHost      string
		GCEConfigFilePath  string
		KubeConfigFilePath string
		PrometheusAddress  string
	}{}
)

// Register registers flags with command line parser
func Register() {
	flag.StringVar(&F.APIServerHost, "apiserver-host", "",
		`The address of the Kubernetes Apiserver to connect to in the format of
protocol://address:port, e.g., http://localhost:8080. If not specified, the
assumption is that the binary runs inside a Kubernetes cluster and local
discovery is attempted.`)
	flag.StringVar(&F.GCEConfigFilePath, "gce-config-file-path", "",
		"Path to a file containing the gce config.")
	flag.StringVar(&F.KubeConfigFilePath, "kube-config-file-path", "",
		"Path to kubeconfig file with authorization and master location information.")
	flag.StringVar(&F.PrometheusAddress, "prometheus-address", ":8944",
		"The address to expose Prometheus metrics")

	flag.Parse()
}
<commit_msg>Change metrics port number to avoid conflicts with VPA<commit_after>/*
Copyright 2018 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    https://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 flags defines global controller flags
package flags

import (
	"flag"
)

var (
	F = struct {
		APIServerHost      string
		GCEConfigFilePath  string
		KubeConfigFilePath string
		PrometheusAddress  string
	}{}
)

// Register registers flags with command line parser
func Register() {
	flag.StringVar(&F.APIServerHost, "apiserver-host", "",
		`The address of the Kubernetes Apiserver to connect to in the format of
protocol://address:port, e.g., http://localhost:8080. If not specified, the
assumption is that the binary runs inside a Kubernetes cluster and local
discovery is attempted.`)
	flag.StringVar(&F.GCEConfigFilePath, "gce-config-file-path", "",
		"Path to a file containing the gce config.")
	flag.StringVar(&F.KubeConfigFilePath, "kube-config-file-path", "",
		"Path to kubeconfig file with authorization and master location information.")
	flag.StringVar(&F.PrometheusAddress, "prometheus-address", ":8910",
		"The address to expose Prometheus metrics")

	flag.Parse()
}
<|endoftext|>


<commit_before>/*
Copyright 2011 Google 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,
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 server

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os/user"
	"time"

	"camlistore.org/pkg/auth"
	"camlistore.org/pkg/blobserver"
	"camlistore.org/pkg/jsonconfig"
	"camlistore.org/pkg/search"
)

// RootHandler handles serving the about/splash page.
type RootHandler struct {
	// Stealth determines whether we hide from non-authenticated
	// clients.
	Stealth bool

	OwnerName string // for display purposes only

	// URL prefixes (path or full URL) to the primary blob and
	// search root.
	BlobRoot   string
	SearchRoot string

	Storage blobserver.Storage // of BlobRoot, or nil
	Search  *search.Handler    // of SearchRoot, or nil

	ui *UIHandler // or nil, if none configured
}

func init() {
	blobserver.RegisterHandlerConstructor("root", newRootFromConfig)
}

func newRootFromConfig(ld blobserver.Loader, conf jsonconfig.Obj) (h http.Handler, err error) {
	u, err := user.Current()
	if err != nil {
		return
	}
	root := &RootHandler{
		BlobRoot:   conf.OptionalString("blobRoot", ""),
		SearchRoot: conf.OptionalString("searchRoot", ""),
		OwnerName:  conf.OptionalString("ownerName", u.Name),
	}
	root.Stealth = conf.OptionalBool("stealth", false)
	if err = conf.Validate(); err != nil {
		return
	}

	if root.BlobRoot != "" {
		bs, err := ld.GetStorage(root.BlobRoot)
		if err != nil {
			return nil, fmt.Errorf("Root handler's blobRoot of %q error: %v", root.BlobRoot, err)
		}
		root.Storage = bs
	}

	if root.SearchRoot != "" {
		h, _ := ld.GetHandler(root.SearchRoot)
		root.Search = h.(*search.Handler)
	}

	return root, nil
}

func (rh *RootHandler) registerUIHandler(h *UIHandler) {
	rh.ui = h
}

func (rh *RootHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if wantsDiscovery(req) {
		if auth.IsAuthorized(req) {
			rh.serveDiscovery(rw, req)
			return
		}
		if !rh.Stealth {
			http.Error(rw, "Unauthorized", http.StatusUnauthorized)
		}
		return
	}

	if rh.Stealth {
		return
	}

	configLink := ""
	if auth.LocalhostAuthorized(req) {
		configLink = "
If you're coming from localhost, hit /setup.
"
	}
	fmt.Fprintf(rw, "This is camlistored, a "+
		"Camlistore server."+
		"%s\n", configLink)
}

func (rh *RootHandler) serveDiscovery(rw http.ResponseWriter, req *http.Request) {
	m := map[string]interface{}{
		"blobRoot":   rh.BlobRoot,
		"searchRoot": rh.SearchRoot,
		"ownerName":  rh.OwnerName,
	}
	if gener, ok := rh.Storage.(blobserver.Generationer); ok {
		initTime, gen, err := gener.StorageGeneration()
		if err != nil {
			m["storageGenerationError"] = err.Error()
		} else {
			m["storageInitTime"] = initTime.UTC().Format(time.RFC3339)
			m["storageGeneration"] = gen
		}
	}
	if rh.ui != nil {
		rh.ui.populateDiscoveryMap(m)
	}
	discoveryHelper(rw, req, m)
}

func discoveryHelper(rw http.ResponseWriter, req *http.Request, m map[string]interface{}) {
	rw.Header().Set("Content-Type", "text/javascript")
	if cb := req.FormValue("cb"); identPattern.MatchString(cb) {
		fmt.Fprintf(rw, "%s(", cb)
		defer rw.Write([]byte(");\n"))
	} else if v := req.FormValue("var"); identOrDotPattern.MatchString(v) {
		fmt.Fprintf(rw, "var %s = ", v)
		defer rw.Write([]byte(";\n"))
	}
	bytes, _ := json.MarshalIndent(m, "", "  ")
	rw.Write(bytes)
}
<commit_msg>root discovery: drop var, so things are on window.<commit_after>/*
Copyright 2011 Google 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,
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 server

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os/user"
	"time"

	"camlistore.org/pkg/auth"
	"camlistore.org/pkg/blobserver"
	"camlistore.org/pkg/jsonconfig"
	"camlistore.org/pkg/search"
)

// RootHandler handles serving the about/splash page.
type RootHandler struct {
	// Stealth determines whether we hide from non-authenticated
	// clients.
	Stealth bool

	OwnerName string // for display purposes only

	// URL prefixes (path or full URL) to the primary blob and
	// search root.
	BlobRoot   string
	SearchRoot string

	Storage blobserver.Storage // of BlobRoot, or nil
	Search  *search.Handler    // of SearchRoot, or nil

	ui *UIHandler // or nil, if none configured
}

func init() {
	blobserver.RegisterHandlerConstructor("root", newRootFromConfig)
}

func newRootFromConfig(ld blobserver.Loader, conf jsonconfig.Obj) (h http.Handler, err error) {
	u, err := user.Current()
	if err != nil {
		return
	}
	root := &RootHandler{
		BlobRoot:   conf.OptionalString("blobRoot", ""),
		SearchRoot: conf.OptionalString("searchRoot", ""),
		OwnerName:  conf.OptionalString("ownerName", u.Name),
	}
	root.Stealth = conf.OptionalBool("stealth", false)
	if err = conf.Validate(); err != nil {
		return
	}

	if root.BlobRoot != "" {
		bs, err := ld.GetStorage(root.BlobRoot)
		if err != nil {
			return nil, fmt.Errorf("Root handler's blobRoot of %q error: %v", root.BlobRoot, err)
		}
		root.Storage = bs
	}

	if root.SearchRoot != "" {
		h, _ := ld.GetHandler(root.SearchRoot)
		root.Search = h.(*search.Handler)
	}

	return root, nil
}

func (rh *RootHandler) registerUIHandler(h *UIHandler) {
	rh.ui = h
}

func (rh *RootHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if wantsDiscovery(req) {
		if auth.IsAuthorized(req) {
			rh.serveDiscovery(rw, req)
			return
		}
		if !rh.Stealth {
			http.Error(rw, "Unauthorized", http.StatusUnauthorized)
		}
		return
	}

	if rh.Stealth {
		return
	}

	configLink := ""
	if auth.LocalhostAuthorized(req) {
		configLink = "
If you're coming from localhost, hit /setup.
"
	}
	fmt.Fprintf(rw, "This is camlistored, a "+
		"Camlistore server."+
		"%s\n", configLink)
}

func (rh *RootHandler) serveDiscovery(rw http.ResponseWriter, req *http.Request) {
	m := map[string]interface{}{
		"blobRoot":   rh.BlobRoot,
		"searchRoot": rh.SearchRoot,
		"ownerName":  rh.OwnerName,
	}
	if gener, ok := rh.Storage.(blobserver.Generationer); ok {
		initTime, gen, err := gener.StorageGeneration()
		if err != nil {
			m["storageGenerationError"] = err.Error()
		} else {
			m["storageInitTime"] = initTime.UTC().Format(time.RFC3339)
			m["storageGeneration"] = gen
		}
	}
	if rh.ui != nil {
		rh.ui.populateDiscoveryMap(m)
	}
	discoveryHelper(rw, req, m)
}

func discoveryHelper(rw http.ResponseWriter, req *http.Request, m map[string]interface{}) {
	rw.Header().Set("Content-Type", "text/javascript")
	if cb := req.FormValue("cb"); identPattern.MatchString(cb) {
		fmt.Fprintf(rw, "%s(", cb)
		defer rw.Write([]byte(");\n"))
	} else if v := req.FormValue("var"); identOrDotPattern.MatchString(v) {
		fmt.Fprintf(rw, "%s = ", v)
		defer rw.Write([]byte(";\n"))
	}
	bytes, _ := json.MarshalIndent(m, "", "  ")
	rw.Write(bytes)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package util

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"io/ioutil"
	"math/big"
	"net"
	"os"
	"path/filepath"
	"time"

	"github.com/pkg/errors"
)

func GenerateCACert(certPath, keyPath string) error {
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return errors.Wrap(err, "Error generating rsa key")
	}

	template := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			CommonName: "minikubeCA",
		},
		NotBefore: time.Now(),
		NotAfter:  time.Now().Add(time.Hour * 24 * 365 * 10),

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		IsCA: true,
	}

	return writeCertsAndKeys(&template, certPath, priv, keyPath, &template, priv)
}

// You may also specify additional subject alt names (either ip or dns names) for the certificate
// The certificate will be created with file mode 0644. The key will be created with file mode 0600.
// If the certificate or key files already exist, they will be overwritten.
// Any parent directories of the certPath or keyPath will be created as needed with file mode 0755.
func GenerateSignedCert(certPath, keyPath string, ips []net.IP, alternateDNS []string, signerCertPath, signerKeyPath string) error {
	signerCertBytes, err := ioutil.ReadFile(signerCertPath)
	if err != nil {
		return errors.Wrap(err, "Error reading file: signerCertPath")
	}
	decodedSignerCert, _ := pem.Decode(signerCertBytes)
	if decodedSignerCert == nil {
		return errors.New("Unable to decode certificate.")
	}
	signerCert, err := x509.ParseCertificate(decodedSignerCert.Bytes)
	if err != nil {
		return errors.Wrap(err, "Error parsing certificate: decodedSignerCert.Bytes")
	}
	signerKeyBytes, err := ioutil.ReadFile(signerKeyPath)
	if err != nil {
		return errors.Wrap(err, "Error reading file: signerKeyPath")
	}
	decodedSignerKey, _ := pem.Decode(signerKeyBytes)
	if decodedSignerKey == nil {
		return errors.New("Unable to decode key.")
	}
	signerKey, err := x509.ParsePKCS1PrivateKey(decodedSignerKey.Bytes)
	if err != nil {
		return errors.Wrap(err, "Error parsing prive key: decodedSignerKey.Bytes")
	}

	template := x509.Certificate{
		SerialNumber: big.NewInt(2),
		Subject: pkix.Name{
			CommonName: "minikube",
		},
		NotBefore: time.Now(),
		NotAfter:  time.Now().Add(time.Hour * 24 * 365),

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
	}

	template.IPAddresses = append(template.IPAddresses, ips...)
	template.DNSNames = append(template.DNSNames, alternateDNS...)

	priv, err := loadOrGeneratePrivateKey(keyPath)
	if err != nil {
		return errors.Wrap(err, "Error loading or generating private key: keyPath")
	}

	return writeCertsAndKeys(&template, certPath, priv, keyPath, signerCert, signerKey)
}

func loadOrGeneratePrivateKey(keyPath string) (*rsa.PrivateKey, error) {
	keyBytes, err := ioutil.ReadFile(keyPath)
	if err == nil {
		decodedKey, _ := pem.Decode(keyBytes)
		if decodedKey != nil {
			priv, err := x509.ParsePKCS1PrivateKey(decodedKey.Bytes)
			if err == nil {
				return priv, nil
			}
		}
	}
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, errors.Wrap(err, "Error generating RSA key")
	}
	return priv, nil
}

func writeCertsAndKeys(template *x509.Certificate, certPath string, signeeKey *rsa.PrivateKey, keyPath string, parent *x509.Certificate, signingKey *rsa.PrivateKey) error {
	derBytes, err := x509.CreateCertificate(rand.Reader, template, parent, &signeeKey.PublicKey, signingKey)
	if err != nil {
		return errors.Wrap(err, "Error creating certificate")
	}

	certBuffer := bytes.Buffer{}
	if err := pem.Encode(&certBuffer, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes}); err != nil {
		return errors.Wrap(err, "Error encoding certificate")
	}

	keyBuffer := bytes.Buffer{}
	if err := pem.Encode(&keyBuffer, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(signeeKey)}); err != nil {
		return errors.Wrap(err, "Error encoding key")
	}

	if err := os.MkdirAll(filepath.Dir(certPath), os.FileMode(0755)); err != nil {
		return errors.Wrap(err, "Error creating certificate directory")
	}
	if err := ioutil.WriteFile(certPath, certBuffer.Bytes(), os.FileMode(0644)); err != nil {
		return errors.Wrap(err, "Error writing certificate to cert path")
	}

	if err := os.MkdirAll(filepath.Dir(keyPath), os.FileMode(0755)); err != nil {
		return errors.Wrap(err, "Error creating key directory")
	}
	if err := ioutil.WriteFile(keyPath, keyBuffer.Bytes(), os.FileMode(0600)); err != nil {
		return errors.Wrap(err, "Error writing key file")
	}

	return nil
}
<commit_msg>Add the 'minikube' user to the 'system:masters' group<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package util

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"io/ioutil"
	"math/big"
	"net"
	"os"
	"path/filepath"
	"time"

	"github.com/pkg/errors"
)

func GenerateCACert(certPath, keyPath string) error {
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return errors.Wrap(err, "Error generating rsa key")
	}

	template := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			CommonName: "minikubeCA",
		},
		NotBefore: time.Now(),
		NotAfter:  time.Now().Add(time.Hour * 24 * 365 * 10),

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		IsCA: true,
	}

	return writeCertsAndKeys(&template, certPath, priv, keyPath, &template, priv)
}

// You may also specify additional subject alt names (either ip or dns names) for the certificate
// The certificate will be created with file mode 0644. The key will be created with file mode 0600.
// If the certificate or key files already exist, they will be overwritten.
// Any parent directories of the certPath or keyPath will be created as needed with file mode 0755.
func GenerateSignedCert(certPath, keyPath string, ips []net.IP, alternateDNS []string, signerCertPath, signerKeyPath string) error {
	signerCertBytes, err := ioutil.ReadFile(signerCertPath)
	if err != nil {
		return errors.Wrap(err, "Error reading file: signerCertPath")
	}
	decodedSignerCert, _ := pem.Decode(signerCertBytes)
	if decodedSignerCert == nil {
		return errors.New("Unable to decode certificate.")
	}
	signerCert, err := x509.ParseCertificate(decodedSignerCert.Bytes)
	if err != nil {
		return errors.Wrap(err, "Error parsing certificate: decodedSignerCert.Bytes")
	}
	signerKeyBytes, err := ioutil.ReadFile(signerKeyPath)
	if err != nil {
		return errors.Wrap(err, "Error reading file: signerKeyPath")
	}
	decodedSignerKey, _ := pem.Decode(signerKeyBytes)
	if decodedSignerKey == nil {
		return errors.New("Unable to decode key.")
	}
	signerKey, err := x509.ParsePKCS1PrivateKey(decodedSignerKey.Bytes)
	if err != nil {
		return errors.Wrap(err, "Error parsing prive key: decodedSignerKey.Bytes")
	}

	template := x509.Certificate{
		SerialNumber: big.NewInt(2),
		Subject: pkix.Name{
			CommonName:   "minikube",
			Organization: []string{"system:masters"},
		},
		NotBefore: time.Now(),
		NotAfter:  time.Now().Add(time.Hour * 24 * 365),

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
	}

	template.IPAddresses = append(template.IPAddresses, ips...)
	template.DNSNames = append(template.DNSNames, alternateDNS...)

	priv, err := loadOrGeneratePrivateKey(keyPath)
	if err != nil {
		return errors.Wrap(err, "Error loading or generating private key: keyPath")
	}

	return writeCertsAndKeys(&template, certPath, priv, keyPath, signerCert, signerKey)
}

func loadOrGeneratePrivateKey(keyPath string) (*rsa.PrivateKey, error) {
	keyBytes, err := ioutil.ReadFile(keyPath)
	if err == nil {
		decodedKey, _ := pem.Decode(keyBytes)
		if decodedKey != nil {
			priv, err := x509.ParsePKCS1PrivateKey(decodedKey.Bytes)
			if err == nil {
				return priv, nil
			}
		}
	}
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, errors.Wrap(err, "Error generating RSA key")
	}
	return priv, nil
}

func writeCertsAndKeys(template *x509.Certificate, certPath string, signeeKey *rsa.PrivateKey, keyPath string, parent *x509.Certificate, signingKey *rsa.PrivateKey) error {
	derBytes, err := x509.CreateCertificate(rand.Reader, template, parent, &signeeKey.PublicKey, signingKey)
	if err != nil {
		return errors.Wrap(err, "Error creating certificate")
	}

	certBuffer := bytes.Buffer{}
	if err := pem.Encode(&certBuffer, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes}); err != nil {
		return errors.Wrap(err, "Error encoding certificate")
	}

	keyBuffer := bytes.Buffer{}
	if err := pem.Encode(&keyBuffer, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(signeeKey)}); err != nil {
		return errors.Wrap(err, "Error encoding key")
	}

	if err := os.MkdirAll(filepath.Dir(certPath), os.FileMode(0755)); err != nil {
		return errors.Wrap(err, "Error creating certificate directory")
	}
	if err := ioutil.WriteFile(certPath, certBuffer.Bytes(), os.FileMode(0644)); err != nil {
		return errors.Wrap(err, "Error writing certificate to cert path")
	}

	if err := os.MkdirAll(filepath.Dir(keyPath), os.FileMode(0755)); err != nil {
		return errors.Wrap(err, "Error creating key directory")
	}
	if err := ioutil.WriteFile(keyPath, keyBuffer.Bytes(), os.FileMode(0600)); err != nil {
		return errors.Wrap(err, "Error writing key file")
	}

	return nil
}
<|endoftext|>


<commit_before>// PromHouse
// Copyright (C) 2017 Percona LLC
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

package storages

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
	"sync"
	"time"

	_ "github.com/kshvakov/clickhouse" // register SQL driver
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/prompb"
	"github.com/sirupsen/logrus"

	"github.com/Percona-Lab/PromHouse/util"
)

const (
	namespace = "promhouse"
	subsystem = "clickhouse"
)

// ClickHouse implements storage interface for the ClickHouse.
type ClickHouse struct {
	db       *sql.DB
	l        *logrus.Entry
	database string

	metrics   map[model.Fingerprint]model.Metric
	metricsRW sync.RWMutex

	mMetricsCurrent prometheus.Gauge
	mSamplesCurrent prometheus.Gauge

	mReads      prometheus.Summary
	mReadErrors *prometheus.CounterVec

	mWrites         prometheus.Summary
	mWriteErrors    *prometheus.CounterVec
	mWrittenLabels  *prometheus.CounterVec
	mWrittenMetrics prometheus.Counter
	mWrittenSamples prometheus.Counter
}

func NewClickHouse(dsn string, database string, init bool) (*ClickHouse, error) {
	l := logrus.WithField("component", "clickhouse")

	db, err := sql.Open("clickhouse", dsn)
	if err != nil {
		return nil, err
	}
	if err = db.Ping(); err != nil {
		return nil, err
	}

	var queries []string
	if init {
		queries = append(queries, fmt.Sprintf(`DROP DATABASE IF EXISTS %s`, database))
	}
	queries = append(queries, fmt.Sprintf(`CREATE DATABASE IF NOT EXISTS %s`, database))

	// TODO use GraphiteMergeTree?
	// TODO remove fingerprint?
	// TODO add __name__, instance, job as separate columns?
	// labels := []string{"instance", "job"}
	queries = append(queries, fmt.Sprintf(`
		CREATE TABLE IF NOT EXISTS %s.samples (
			__date__ Date,
			__labels__ String,
			__fingerprint__ UInt64,
			__ts__ Int64,
			__value__ Float64
		)
		ENGINE = MergeTree(__date__, (__fingerprint__, __ts__), 8192)`, database))

	queries = append(queries, fmt.Sprintf(`
		CREATE TABLE IF NOT EXISTS %s.metrics (
			date Date,
			fingerprint UInt64,
			labels String
		)
		ENGINE = MergeTree(date, fingerprint, 8192)`, database))

	for _, q := range queries {
		l.Infof("Executing: %s", q)
		if _, err = db.Exec(q); err != nil {
			return nil, err
		}
	}

	ch := &ClickHouse{
		db:       db,
		l:        l,
		database: database,

		metrics: make(map[model.Fingerprint]model.Metric, 8192),

		mMetricsCurrent: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "metrics_current",
			Help:      "Current number of stored metrics.",
		}),
		mSamplesCurrent: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "samples_current",
			Help:      "Current number of stored samples.",
		}),

		mReads: prometheus.NewSummary(prometheus.SummaryOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "reads",
			Help:      "Durations of successful reads.",
		}),
		mReadErrors: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "read_errors",
			Help:      "Number of read errors by type: canceled, other.",
		}, []string{"type"}),

		mWrites: prometheus.NewSummary(prometheus.SummaryOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "writes",
			Help:      "Durations of committed writes.",
		}),
		mWriteErrors: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "write_errors",
			Help:      "Number of write errors by type: canceled, other.",
		}, []string{"type"}),
		mWrittenLabels: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "written_labels",
			Help:      "Number of written labels by name.",
		}, []string{"name"}),
		mWrittenMetrics: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "written_metrics",
			Help:      "Number of written metrics.",
		}),
		mWrittenSamples: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "written_samples",
			Help:      "Number of written samples.",
		}),
	}

	ch.mReadErrors.WithLabelValues("canceled").Set(0)
	ch.mReadErrors.WithLabelValues("other").Set(0)
	ch.mWriteErrors.WithLabelValues("canceled").Set(0)
	ch.mWriteErrors.WithLabelValues("other").Set(0)

	return ch, nil
}

// makeMetric converts a slice of labels to a metric.
func makeMetric(labels []*prompb.Label) model.Metric {
	metric := make(model.Metric, len(labels))
	for _, l := range labels {
		metric[model.LabelName(l.Name)] = model.LabelValue(l.Value)
	}
	return metric
}

func (ch *ClickHouse) Describe(c chan<- *prometheus.Desc) {
	ch.mMetricsCurrent.Describe(c)
	ch.mSamplesCurrent.Describe(c)

	ch.mReads.Describe(c)
	ch.mReadErrors.Describe(c)

	ch.mWrites.Describe(c)
	ch.mWriteErrors.Describe(c)
	ch.mWrittenLabels.Describe(c)
	ch.mWrittenMetrics.Describe(c)
	ch.mWrittenSamples.Describe(c)
}

func (ch *ClickHouse) Collect(c chan<- prometheus.Metric) {
	var count uint64
	query := fmt.Sprintf("SELECT COUNT(*) FROM %s.metrics", ch.database)
	if err := ch.db.QueryRow(query).Scan(&count); err != nil {
		ch.l.Error(err)
		return
	}
	ch.mMetricsCurrent.Set(float64(count))

	query = fmt.Sprintf("SELECT COUNT(*) FROM %s.samples", ch.database)
	if err := ch.db.QueryRow(query).Scan(&count); err != nil {
		ch.l.Error(err)
		return
	}
	ch.mSamplesCurrent.Set(float64(count))

	ch.mMetricsCurrent.Collect(c)
	ch.mSamplesCurrent.Collect(c)

	ch.mReads.Collect(c)
	ch.mReadErrors.Collect(c)

	ch.mWrites.Collect(c)
	ch.mWriteErrors.Collect(c)
	ch.mWrittenLabels.Collect(c)
	ch.mWrittenMetrics.Collect(c)
	ch.mWrittenSamples.Collect(c)
}

func (ch *ClickHouse) Read(ctx context.Context, queries []Query) (res *prompb.ReadResponse, err error) {
	start := time.Now()
	defer func() {
		if err == nil {
			ch.mReads.Observe(time.Since(start).Seconds())
			return
		}

		t := "other"
		if err == context.Canceled {
			t = "canceled"
		}
		ch.mReadErrors.WithLabelValues(t).Inc()
	}()

	res = &prompb.ReadResponse{
		Results: make([]*prompb.QueryResult, len(queries)),
	}
	for i, q := range queries {
		res.Results[i] = new(prompb.QueryResult)
		query := fmt.Sprintf("SELECT __labels__, __fingerprint__, __ts__, __value__ FROM %s.samples WHERE __ts__ >= ? AND __ts__ <= ?", ch.database)
		args := []interface{}{int64(q.Start), int64(q.End)}
		for _, matcher := range q.Matchers {
			switch matcher.Type {
			case MatchEqual:
				query += fmt.Sprintf(" AND visitParamExtractString(__labels__, '%s') = ?", matcher.Name)
				args = append(args, matcher.Value)
			case MatchNotEqual:
				query += fmt.Sprintf(" AND visitParamExtractString(__labels__, '%s') != ?", matcher.Name)
				args = append(args, matcher.Value)
			case MatchRegexp:
				query += fmt.Sprintf(" AND match(visitParamExtractString(__labels__, '%s'), ?)", matcher.Name)
				args = append(args, "^(?:"+matcher.Value+")$")
			case MatchNotRegexp:
				query += fmt.Sprintf(" AND NOT match(visitParamExtractString(__labels__, '%s'), ?)", matcher.Name)
				args = append(args, "^(?:"+matcher.Value+")$")
			default:
				ch.l.Panicf("unexpected match type: %d", matcher.Type)
			}
		}
		query += " ORDER BY __fingerprint__, __ts__"

		var rows *sql.Rows
		rows, err = ch.db.Query(query, args...)
		if err != nil {
			return
		}

		var ts *prompb.TimeSeries
		var b []byte
		var fingerprint, prevFingerprint uint64
		var t int64
		var value float64
		for rows.Next() {
			if err = rows.Scan(&b, &fingerprint, &t, &value); err != nil {
				return
			}
			if fingerprint != prevFingerprint {
				prevFingerprint = fingerprint
				if ts != nil {
					res.Results[i].Timeseries = append(res.Results[i].Timeseries, ts)
				}

				var ls model.LabelSet
				if err = ls.UnmarshalJSON(b); err != nil {
					return
				}
				labels := make([]*prompb.Label, 0, len(ls))
				for n, v := range ls {
					labels = append(labels, &prompb.Label{
						Name:  string(n),
						Value: string(v),
					})
				}
				ts = &prompb.TimeSeries{
					Labels: labels,
				}
			}
			ts.Samples = append(ts.Samples, &prompb.Sample{
				Timestamp: t,
				Value:     value,
			})
		}
		if err = rows.Close(); err != nil {
			return
		}
		if ts != nil {
			res.Results[i].Timeseries = append(res.Results[i].Timeseries, ts)
		}
	}

	return
}

func inTransaction(ctx context.Context, db *sql.DB, f func(*sql.Tx) error) (err error) {
	var tx *sql.Tx
	if tx, err = db.BeginTx(ctx, nil); err != nil {
		return
	}
	defer func() {
		if err == nil {
			err = tx.Commit()
		} else {
			tx.Rollback()
		}
	}()
	err = f(tx)
	return
}

func (ch *ClickHouse) Write(ctx context.Context, data *prompb.WriteRequest) (err error) {
	start := time.Now()
	defer func() {
		if err == nil {
			ch.mWrites.Observe(time.Since(start).Seconds())
			return
		}

		t := "other"
		if err == context.Canceled {
			t = "canceled"
		}
		ch.mWriteErrors.WithLabelValues(t).Inc()
	}()

	// make metrics and calculate fingerprints
	metrics := make(map[model.Fingerprint]model.Metric, len(data.Timeseries))
	for _, ts := range data.Timeseries {
		m := makeMetric(ts.Labels)
		metrics[m.Fingerprint()] = m
	}

	// find new metrics
	newMetrics := make([]model.Fingerprint, 0, len(metrics))
	ch.metricsRW.Lock()
	for f, m := range metrics {
		_, ok := ch.metrics[f]
		if !ok {
			newMetrics = append(newMetrics, f)
			ch.metrics[f] = m
		}
	}
	ch.metricsRW.Unlock()

	// write new metrics
	err = inTransaction(ctx, ch.db, func(tx *sql.Tx) error {
		placeholders := strings.Repeat("?, ", 3)
		query := fmt.Sprintf(
			`INSERT INTO %s.metrics (date, fingerprint, labels) VALUES (%s)`,
			ch.database, placeholders[:len(placeholders)-2], // cut last ", "
		)
		var stmt *sql.Stmt
		if stmt, err = tx.PrepareContext(ctx, query); err != nil {
			return err
		}

		args := make([]interface{}, 3)
		args[0] = model.Now().Time()
		for _, f := range newMetrics {
			args[1] = uint64(f)
			args[2] = util.MarshalMetric(metrics[f])
			if _, err = stmt.ExecContext(ctx, args...); err != nil {
				return err
			}
		}

		return stmt.Close()
	})
	if err != nil {
		return
	}

	// write samples
	var samples int
	err = inTransaction(ctx, ch.db, func(tx *sql.Tx) error {
		placeholders := strings.Repeat("?, ", 5)
		query := fmt.Sprintf(
			`INSERT INTO %s.samples (__labels__, __fingerprint__, __ts__, __value__, __date__) VALUES (%s)`,
			ch.database, placeholders[:len(placeholders)-2], // cut last ", "
		)
		var stmt *sql.Stmt
		if stmt, err = tx.PrepareContext(ctx, query); err != nil {
			return err
		}

		args := make([]interface{}, 5)
		for _, ts := range data.Timeseries {
			m := makeMetric(ts.Labels)
			args[0] = util.MarshalMetric(m)
			args[1] = uint64(m.Fingerprint())

			for _, s := range ts.Samples {
				args[len(args)-3] = s.Timestamp
				args[len(args)-2] = s.Value
				args[len(args)-1] = model.Time(s.Timestamp).Time()
				if _, err = stmt.ExecContext(ctx, args...); err != nil {
					return err
				}
				samples++
			}
		}

		return stmt.Close()
	})
	if err != nil {
		return
	}

	ch.mWrittenMetrics.Add(float64(len(newMetrics)))
	ch.mWrittenSamples.Add(float64(samples))
	ch.l.Debugf("Wrote %s new metrics, %d samples.", len(newMetrics), samples)
	return
}

// check interface
var _ Storage = (*ClickHouse)(nil)
<commit_msg>Load existing metrics during startup.<commit_after>// PromHouse
// Copyright (C) 2017 Percona LLC
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

package storages

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
	"sync"
	"time"

	_ "github.com/kshvakov/clickhouse" // register SQL driver
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/prompb"
	"github.com/sirupsen/logrus"

	"github.com/Percona-Lab/PromHouse/util"
)

const (
	namespace = "promhouse"
	subsystem = "clickhouse"
)

// ClickHouse implements storage interface for the ClickHouse.
type ClickHouse struct {
	db       *sql.DB
	l        *logrus.Entry
	database string

	metrics   map[model.Fingerprint]model.Metric
	metricsRW sync.RWMutex

	mMetricsCurrent prometheus.Gauge
	mSamplesCurrent prometheus.Gauge

	mReads      prometheus.Summary
	mReadErrors *prometheus.CounterVec

	mWrites         prometheus.Summary
	mWriteErrors    *prometheus.CounterVec
	mWrittenLabels  *prometheus.CounterVec
	mWrittenMetrics prometheus.Counter
	mWrittenSamples prometheus.Counter
}

func NewClickHouse(dsn string, database string, init bool) (*ClickHouse, error) {
	l := logrus.WithField("component", "clickhouse")

	db, err := sql.Open("clickhouse", dsn)
	if err != nil {
		return nil, err
	}
	if err = db.Ping(); err != nil {
		return nil, err
	}

	var queries []string
	if init {
		queries = append(queries, fmt.Sprintf(`DROP DATABASE IF EXISTS %s`, database))
	}
	queries = append(queries, fmt.Sprintf(`CREATE DATABASE IF NOT EXISTS %s`, database))

	// TODO use GraphiteMergeTree?
	// TODO remove fingerprint?
	// TODO add __name__, instance, job as separate columns?
	// labels := []string{"instance", "job"}
	queries = append(queries, fmt.Sprintf(`
		CREATE TABLE IF NOT EXISTS %s.samples (
			__date__ Date,
			__labels__ String,
			__fingerprint__ UInt64,
			__ts__ Int64,
			__value__ Float64
		)
		ENGINE = MergeTree(__date__, (__fingerprint__, __ts__), 8192)`, database))

	queries = append(queries, fmt.Sprintf(`
		CREATE TABLE IF NOT EXISTS %s.metrics (
			date Date,
			fingerprint UInt64,
			labels String
		)
		ENGINE = MergeTree(date, fingerprint, 8192)`, database))

	for _, q := range queries {
		l.Infof("Executing: %s", q)
		if _, err = db.Exec(q); err != nil {
			return nil, err
		}
	}

	// load existing metrics
	metrics := make(map[model.Fingerprint]model.Metric, 8192)
	q := fmt.Sprintf(`SELECT fingerprint, labels FROM %s.metrics`, database)
	rows, err := db.Query(q)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var f uint64
	var b []byte
	for rows.Next() {
		if err = rows.Scan(&f, &b); err != nil {
			return nil, err
		}
		var ls model.LabelSet
		if err = ls.UnmarshalJSON(b); err != nil {
			return nil, err
		}
		metrics[model.Fingerprint(f)] = model.Metric(ls)
	}
	if err = rows.Err(); err != nil {
		return nil, err
	}
	l.Debugf("Loaded %d existing metrics.", len(metrics))

	ch := &ClickHouse{
		db:       db,
		l:        l,
		database: database,

		metrics: metrics,

		mMetricsCurrent: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "metrics_current",
			Help:      "Current number of stored metrics.",
		}),
		mSamplesCurrent: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "samples_current",
			Help:      "Current number of stored samples.",
		}),

		mReads: prometheus.NewSummary(prometheus.SummaryOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "reads",
			Help:      "Durations of successful reads.",
		}),
		mReadErrors: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "read_errors",
			Help:      "Number of read errors by type: canceled, other.",
		}, []string{"type"}),

		mWrites: prometheus.NewSummary(prometheus.SummaryOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "writes",
			Help:      "Durations of committed writes.",
		}),
		mWriteErrors: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "write_errors",
			Help:      "Number of write errors by type: canceled, other.",
		}, []string{"type"}),
		mWrittenLabels: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "written_labels",
			Help:      "Number of written labels by name.",
		}, []string{"name"}),
		mWrittenMetrics: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "written_metrics",
			Help:      "Number of written metrics.",
		}),
		mWrittenSamples: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "written_samples",
			Help:      "Number of written samples.",
		}),
	}

	ch.mReadErrors.WithLabelValues("canceled").Set(0)
	ch.mReadErrors.WithLabelValues("other").Set(0)
	ch.mWriteErrors.WithLabelValues("canceled").Set(0)
	ch.mWriteErrors.WithLabelValues("other").Set(0)

	return ch, nil
}

// makeMetric converts a slice of labels to a metric.
func makeMetric(labels []*prompb.Label) model.Metric {
	metric := make(model.Metric, len(labels))
	for _, l := range labels {
		metric[model.LabelName(l.Name)] = model.LabelValue(l.Value)
	}
	return metric
}

func (ch *ClickHouse) Describe(c chan<- *prometheus.Desc) {
	ch.mMetricsCurrent.Describe(c)
	ch.mSamplesCurrent.Describe(c)

	ch.mReads.Describe(c)
	ch.mReadErrors.Describe(c)

	ch.mWrites.Describe(c)
	ch.mWriteErrors.Describe(c)
	ch.mWrittenLabels.Describe(c)
	ch.mWrittenMetrics.Describe(c)
	ch.mWrittenSamples.Describe(c)
}

func (ch *ClickHouse) Collect(c chan<- prometheus.Metric) {
	var count uint64
	query := fmt.Sprintf("SELECT COUNT(*) FROM %s.metrics", ch.database)
	if err := ch.db.QueryRow(query).Scan(&count); err != nil {
		ch.l.Error(err)
		return
	}
	ch.mMetricsCurrent.Set(float64(count))

	query = fmt.Sprintf("SELECT COUNT(*) FROM %s.samples", ch.database)
	if err := ch.db.QueryRow(query).Scan(&count); err != nil {
		ch.l.Error(err)
		return
	}
	ch.mSamplesCurrent.Set(float64(count))

	ch.mMetricsCurrent.Collect(c)
	ch.mSamplesCurrent.Collect(c)

	ch.mReads.Collect(c)
	ch.mReadErrors.Collect(c)

	ch.mWrites.Collect(c)
	ch.mWriteErrors.Collect(c)
	ch.mWrittenLabels.Collect(c)
	ch.mWrittenMetrics.Collect(c)
	ch.mWrittenSamples.Collect(c)
}

func (ch *ClickHouse) Read(ctx context.Context, queries []Query) (res *prompb.ReadResponse, err error) {
	start := time.Now()
	defer func() {
		if err == nil {
			ch.mReads.Observe(time.Since(start).Seconds())
			return
		}

		t := "other"
		if err == context.Canceled {
			t = "canceled"
		}
		ch.mReadErrors.WithLabelValues(t).Inc()
	}()

	res = &prompb.ReadResponse{
		Results: make([]*prompb.QueryResult, len(queries)),
	}
	for i, q := range queries {
		res.Results[i] = new(prompb.QueryResult)
		query := fmt.Sprintf("SELECT __labels__, __fingerprint__, __ts__, __value__ FROM %s.samples WHERE __ts__ >= ? AND __ts__ <= ?", ch.database)
		args := []interface{}{int64(q.Start), int64(q.End)}
		for _, matcher := range q.Matchers {
			switch matcher.Type {
			case MatchEqual:
				query += fmt.Sprintf(" AND visitParamExtractString(__labels__, '%s') = ?", matcher.Name)
				args = append(args, matcher.Value)
			case MatchNotEqual:
				query += fmt.Sprintf(" AND visitParamExtractString(__labels__, '%s') != ?", matcher.Name)
				args = append(args, matcher.Value)
			case MatchRegexp:
				query += fmt.Sprintf(" AND match(visitParamExtractString(__labels__, '%s'), ?)", matcher.Name)
				args = append(args, "^(?:"+matcher.Value+")$")
			case MatchNotRegexp:
				query += fmt.Sprintf(" AND NOT match(visitParamExtractString(__labels__, '%s'), ?)", matcher.Name)
				args = append(args, "^(?:"+matcher.Value+")$")
			default:
				ch.l.Panicf("unexpected match type: %d", matcher.Type)
			}
		}
		query += " ORDER BY __fingerprint__, __ts__"

		var rows *sql.Rows
		rows, err = ch.db.Query(query, args...)
		if err != nil {
			return
		}

		var ts *prompb.TimeSeries
		var b []byte
		var fingerprint, prevFingerprint uint64
		var t int64
		var value float64
		for rows.Next() {
			if err = rows.Scan(&b, &fingerprint, &t, &value); err != nil {
				return
			}
			if fingerprint != prevFingerprint {
				prevFingerprint = fingerprint
				if ts != nil {
					res.Results[i].Timeseries = append(res.Results[i].Timeseries, ts)
				}

				var ls model.LabelSet
				if err = ls.UnmarshalJSON(b); err != nil {
					return
				}
				labels := make([]*prompb.Label, 0, len(ls))
				for n, v := range ls {
					labels = append(labels, &prompb.Label{
						Name:  string(n),
						Value: string(v),
					})
				}
				ts = &prompb.TimeSeries{
					Labels: labels,
				}
			}
			ts.Samples = append(ts.Samples, &prompb.Sample{
				Timestamp: t,
				Value:     value,
			})
		}
		if err = rows.Close(); err != nil {
			return
		}
		if ts != nil {
			res.Results[i].Timeseries = append(res.Results[i].Timeseries, ts)
		}
	}

	return
}

func inTransaction(ctx context.Context, db *sql.DB, f func(*sql.Tx) error) (err error) {
	var tx *sql.Tx
	if tx, err = db.BeginTx(ctx, nil); err != nil {
		return
	}
	defer func() {
		if err == nil {
			err = tx.Commit()
		} else {
			tx.Rollback()
		}
	}()
	err = f(tx)
	return
}

func (ch *ClickHouse) Write(ctx context.Context, data *prompb.WriteRequest) (err error) {
	start := time.Now()
	defer func() {
		if err == nil {
			ch.mWrites.Observe(time.Since(start).Seconds())
			return
		}

		t := "other"
		if err == context.Canceled {
			t = "canceled"
		}
		ch.mWriteErrors.WithLabelValues(t).Inc()
	}()

	// make metrics and calculate fingerprints
	metrics := make(map[model.Fingerprint]model.Metric, len(data.Timeseries))
	for _, ts := range data.Timeseries {
		m := makeMetric(ts.Labels)
		metrics[m.Fingerprint()] = m
	}

	// find new metrics
	newMetrics := make([]model.Fingerprint, 0, len(metrics))
	ch.metricsRW.Lock()
	for f, m := range metrics {
		_, ok := ch.metrics[f]
		if !ok {
			newMetrics = append(newMetrics, f)
			ch.metrics[f] = m
		}
	}
	ch.metricsRW.Unlock()

	// write new metrics
	err = inTransaction(ctx, ch.db, func(tx *sql.Tx) error {
		placeholders := strings.Repeat("?, ", 3)
		query := fmt.Sprintf(
			`INSERT INTO %s.metrics (date, fingerprint, labels) VALUES (%s)`,
			ch.database, placeholders[:len(placeholders)-2], // cut last ", "
		)
		var stmt *sql.Stmt
		if stmt, err = tx.PrepareContext(ctx, query); err != nil {
			return err
		}

		args := make([]interface{}, 3)
		args[0] = model.Now().Time()
		for _, f := range newMetrics {
			args[1] = uint64(f)
			args[2] = util.MarshalMetric(metrics[f])
			if _, err = stmt.ExecContext(ctx, args...); err != nil {
				return err
			}
		}

		return stmt.Close()
	})
	if err != nil {
		return
	}

	// write samples
	var samples int
	err = inTransaction(ctx, ch.db, func(tx *sql.Tx) error {
		placeholders := strings.Repeat("?, ", 5)
		query := fmt.Sprintf(
			`INSERT INTO %s.samples (__labels__, __fingerprint__, __ts__, __value__, __date__) VALUES (%s)`,
			ch.database, placeholders[:len(placeholders)-2], // cut last ", "
		)
		var stmt *sql.Stmt
		if stmt, err = tx.PrepareContext(ctx, query); err != nil {
			return err
		}

		args := make([]interface{}, 5)
		for _, ts := range data.Timeseries {
			m := makeMetric(ts.Labels)
			args[0] = util.MarshalMetric(m)
			args[1] = uint64(m.Fingerprint())

			for _, s := range ts.Samples {
				args[len(args)-3] = s.Timestamp
				args[len(args)-2] = s.Value
				args[len(args)-1] = model.Time(s.Timestamp).Time()
				if _, err = stmt.ExecContext(ctx, args...); err != nil {
					return err
				}
				samples++
			}
		}

		return stmt.Close()
	})
	if err != nil {
		return
	}

	ch.mWrittenMetrics.Add(float64(len(newMetrics)))
	ch.mWrittenSamples.Add(float64(samples))
	ch.l.Debugf("Wrote %d new metrics, %d samples.", len(newMetrics), samples)
	return
}

// check interface
var _ Storage = (*ClickHouse)(nil)
<|endoftext|>