<commit_before>package main

import (
	"flag"
	"os"

	_ "net/http/pprof"

	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/sky-uk/feed/controller"
	"github.com/sky-uk/feed/elb"
	"github.com/sky-uk/feed/nginx"
	"github.com/sky-uk/feed/util/cmd"
	"github.com/sky-uk/feed/util/metrics"
)

var (
	debug                          bool
	apiserverURL                   string
	caCertFile                     string
	tokenFile                      string
	clientCertFile                 string
	clientKeyFile                  string
	ingressPort                    int
	ingressAllow                   string
	ingressHealthPort              int
	healthPort                     int
	nginxBinary                    string
	nginxWorkDir                   string
	nginxWorkerProcesses           int
	nginxWorkerConnections         int
	nginxKeepAliveSeconds          int
	nginxBackendKeepalives         int
	nginxBackendKeepaliveSeconds   int
	nginxLogLevel                  string
	nginxTrustedFrontends          string
	nginxServerNamesHashBucketSize int
	nginxServerNamesHashMaxSize    int
	elbLabelValue                  string
	elbRegion                      string
	elbExpectedNumber              int
	pushgatewayURL                 string
	pushgatewayIntervalSeconds     int
	pushgatewayLabels              cmd.KeyValues
)

func init() {
	const (
		defaultAPIServer                      = "https://kubernetes:443"
		defaultCaCertFile                     = "/run/secrets/kubernetes.io/serviceaccount/ca.crt"
		defaultTokenFile                      = ""
		defaultClientCertFile                 = ""
		defaultClientKeyFile                  = ""
		defaultIngressPort                    = 8080
		defaultIngressAllow                   = ""
		defaultIngressHealthPort              = 8081
		defaultHealthPort                     = 12082
		defaultNginxBinary                    = "/usr/sbin/nginx"
		defaultNginxWorkingDir                = "/nginx"
		defaultNginxWorkers                   = 1
		defaultNginxWorkerConnections         = 1024
		defaultNginxKeepAliveSeconds          = 60
		defaultNginxBackendKeepalives         = 512
		defaultNginxBackendKeepaliveSeconds   = 60
		defaultNginxLogLevel                  = "info"
		defaultNginxServerNamesHashBucketSize = -1
		defaultNginxServerNamesHashMaxSize    = -1
		defaultElbLabelValue                  = ""
		defaultElbRegion                      = "eu-west-1"
		defaultElbExpectedNumber              = 0
		defaultPushgatewayIntervalSeconds     = 60
	)

	flag.BoolVar(&debug, "debug", false,
		"Enable debug logging.")
	flag.StringVar(&apiserverURL, "apiserver", defaultAPIServer,
		"Kubernetes API server URL.")
	flag.StringVar(&caCertFile, "cacertfile", defaultCaCertFile,
		"File containing the Kubernetes API server certificate.")
	flag.StringVar(&tokenFile, "tokenfile", defaultTokenFile,
		"File containing kubernetes client authentication token. Leave empty to not use a token.")
	flag.StringVar(&clientCertFile, "client-certfile", defaultClientCertFile,
		"File containing client certificate. Leave empty to not use a client certificate.")
	flag.StringVar(&clientKeyFile, "client-keyfile", defaultClientKeyFile,
		"File containing client key. Leave empty to not use a client certificate.")
	flag.IntVar(&ingressPort, "ingress-port", defaultIngressPort,
		"Port to serve ingress traffic to backend services.")
	flag.IntVar(&ingressHealthPort, "ingress-health-port", defaultIngressHealthPort,
		"Port for ingress /health and /status pages. Should be used by frontends to determine if ingress is available.")
	flag.StringVar(&ingressAllow, "ingress-allow", defaultIngressAllow,
		"Source IP or CIDR to allow ingress access by default. This is overridden by the sky.uk/allow "+
			"annotation on ingress resources. Leave empty to deny all access by default.")
	flag.IntVar(&healthPort, "health-port", defaultHealthPort,
		"Port for checking the health of the ingress controller on /health. Also provides /debug/pprof.")
	flag.StringVar(&nginxBinary, "nginx-binary", defaultNginxBinary,
		"Location of nginx binary.")
	flag.StringVar(&nginxWorkDir, "nginx-workdir", defaultNginxWorkingDir,
		"Directory to store nginx files. Also the location of the nginx.tmpl file.")
	flag.IntVar(&nginxWorkerProcesses, "nginx-workers", defaultNginxWorkers,
		"Number of nginx worker processes.")
	flag.IntVar(&nginxWorkerConnections, "nginx-worker-connections", defaultNginxWorkerConnections,
		"Max number of connections per nginx worker. Includes both client and proxy connections.")
	flag.IntVar(&nginxKeepAliveSeconds, "nginx-keepalive-seconds", defaultNginxKeepAliveSeconds,
		"Keep alive time for persistent client connections to nginx. Should generally be set larger than frontend "+
			"keep alive times to prevent stale connections.")
	flag.IntVar(&nginxBackendKeepalives, "nginx-backend-keepalive-count", defaultNginxBackendKeepalives,
		"Maximum number of keepalive connections per backend service. Keepalive connections count against"+
			" nginx-worker-connections limit, and will be restricted by that global limit as well.")
	flag.IntVar(&nginxBackendKeepaliveSeconds, "nginx-backend-keepalive-seconds", defaultNginxBackendKeepaliveSeconds,
		"Time to keep backend keepalive connections open. This should generally be set smaller than backend service keepalive "+
			"times to prevent stale connections.")
	flag.StringVar(&nginxLogLevel, "nginx-loglevel", defaultNginxLogLevel,
		"Log level for nginx. See http://nginx.org/en/docs/ngx_core_module.html#error_log for levels.")
	flag.IntVar(&nginxServerNamesHashBucketSize, "nginx-server-names-hash-bucket-size", defaultNginxServerNamesHashBucketSize,
		"Sets the bucket size for the server names hash tables. Setting this to 0 or less will exclude this "+
			"config from the nginx conf file. The details of setting up hash tables are provided "+
			"in a separate document. http://nginx.org/en/docs/hash.html")
	flag.IntVar(&nginxServerNamesHashMaxSize, "nginx-server-names-hash-max-size", defaultNginxServerNamesHashMaxSize,
		"Sets the maximum size of the server names hash tables. Setting this to 0 or less will exclude this "+
			"config from the nginx conf file. The details of setting up hash tables are provided "+
			"in a separate document. http://nginx.org/en/docs/hash.html")
	flag.StringVar(&nginxTrustedFrontends, "nginx-trusted-frontends", "",
		"Comma separated list of CIDRs to trust when determining the client's real IP from the "+
			"X-Forwarded-For header. The client IP is used for allowing or denying ingress access. "+
			"This will typically be the ELB subnet.")
	flag.StringVar(&elbLabelValue, "elb-label-value", defaultElbLabelValue,
		"Attach to ELBs tagged with "+elb.ElbTag+"=value. Leave empty to not attach.")
	flag.IntVar(&elbExpectedNumber, "elb-expected-number", defaultElbExpectedNumber,
		"Expected number of ELBs to attach to. If 0 the controller will not check,"+
			" otherwise it fails to start if it can't attach to this number.")
	flag.StringVar(&elbRegion, "elb-region", defaultElbRegion,
		"AWS region for ELBs.")
	flag.StringVar(&pushgatewayURL, "pushgateway", "",
		"Prometheus pushgateway URL for pushing metrics. Leave blank to not push metrics.")
	flag.IntVar(&pushgatewayIntervalSeconds, "pushgateway-interval", defaultPushgatewayIntervalSeconds,
		"Interval in seconds for pushing metrics.")
	flag.Var(&pushgatewayLabels, "pushgateway-label",
		"A label=value pair to attach to metrics pushed to prometheus. Specify multiple times for multiple labels.")
}

func main() {
	flag.Parse()

	cmd.ConfigureLogging(debug)
	cmd.ConfigureMetrics("feed-ingress", pushgatewayLabels, pushgatewayURL, pushgatewayIntervalSeconds)

	client := cmd.CreateK8sClient(caCertFile, tokenFile, apiserverURL, clientCertFile, clientKeyFile)
	updaters := createIngressUpdaters()

	controller := controller.New(controller.Config{
		KubernetesClient: client,
		Updaters:         updaters,
		DefaultAllow:     ingressAllow,
	})

	cmd.AddHealthMetrics(controller, metrics.PrometheusIngressSubsystem)
	cmd.AddHealthPort(controller, healthPort)
	cmd.AddSignalHandler(controller)

	err := controller.Start()
	if err != nil {
		log.Error("Error while starting controller: ", err)
		os.Exit(-1)
	}
	log.Info("Controller started")

	select {}
}

func createIngressUpdaters() []controller.Updater {
	frontend := elb.New(elbRegion, elbLabelValue, elbExpectedNumber)
	trustedFrontends := []string{}
	if nginxTrustedFrontends != "" {
		trustedFrontends = strings.Split(nginxTrustedFrontends, ",")
	}
	proxy := nginx.New(nginx.Conf{
		BinaryLocation:            nginxBinary,
		IngressPort:               ingressPort,
		WorkingDir:                nginxWorkDir,
		WorkerProcesses:           nginxWorkerProcesses,
		WorkerConnections:         nginxWorkerConnections,
		KeepaliveSeconds:          nginxKeepAliveSeconds,
		BackendKeepalives:         nginxBackendKeepalives,
		BackendKeepaliveSeconds:   nginxBackendKeepaliveSeconds,
		ServerNamesHashBucketSize: nginxServerNamesHashBucketSize,
		ServerNamesHashMaxSize:    nginxServerNamesHashMaxSize,
		HealthPort:                ingressHealthPort,
		TrustedFrontends:          trustedFrontends,
	})
	return []controller.Updater{frontend, proxy}
}
<commit_msg>Fix `feed-ingress` to set nginx loglevel<commit_after>package main

import (
	"flag"
	"os"

	_ "net/http/pprof"

	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/sky-uk/feed/controller"
	"github.com/sky-uk/feed/elb"
	"github.com/sky-uk/feed/nginx"
	"github.com/sky-uk/feed/util/cmd"
	"github.com/sky-uk/feed/util/metrics"
)

var (
	debug                          bool
	apiserverURL                   string
	caCertFile                     string
	tokenFile                      string
	clientCertFile                 string
	clientKeyFile                  string
	ingressPort                    int
	ingressAllow                   string
	ingressHealthPort              int
	healthPort                     int
	nginxBinary                    string
	nginxWorkDir                   string
	nginxWorkerProcesses           int
	nginxWorkerConnections         int
	nginxKeepAliveSeconds          int
	nginxBackendKeepalives         int
	nginxBackendKeepaliveSeconds   int
	nginxLogLevel                  string
	nginxTrustedFrontends          string
	nginxServerNamesHashBucketSize int
	nginxServerNamesHashMaxSize    int
	elbLabelValue                  string
	elbRegion                      string
	elbExpectedNumber              int
	pushgatewayURL                 string
	pushgatewayIntervalSeconds     int
	pushgatewayLabels              cmd.KeyValues
)

func init() {
	const (
		defaultAPIServer                      = "https://kubernetes:443"
		defaultCaCertFile                     = "/run/secrets/kubernetes.io/serviceaccount/ca.crt"
		defaultTokenFile                      = ""
		defaultClientCertFile                 = ""
		defaultClientKeyFile                  = ""
		defaultIngressPort                    = 8080
		defaultIngressAllow                   = ""
		defaultIngressHealthPort              = 8081
		defaultHealthPort                     = 12082
		defaultNginxBinary                    = "/usr/sbin/nginx"
		defaultNginxWorkingDir                = "/nginx"
		defaultNginxWorkers                   = 1
		defaultNginxWorkerConnections         = 1024
		defaultNginxKeepAliveSeconds          = 60
		defaultNginxBackendKeepalives         = 512
		defaultNginxBackendKeepaliveSeconds   = 60
		defaultNginxLogLevel                  = "warn"
		defaultNginxServerNamesHashBucketSize = -1
		defaultNginxServerNamesHashMaxSize    = -1
		defaultElbLabelValue                  = ""
		defaultElbRegion                      = "eu-west-1"
		defaultElbExpectedNumber              = 0
		defaultPushgatewayIntervalSeconds     = 60
	)

	flag.BoolVar(&debug, "debug", false,
		"Enable debug logging.")
	flag.StringVar(&apiserverURL, "apiserver", defaultAPIServer,
		"Kubernetes API server URL.")
	flag.StringVar(&caCertFile, "cacertfile", defaultCaCertFile,
		"File containing the Kubernetes API server certificate.")
	flag.StringVar(&tokenFile, "tokenfile", defaultTokenFile,
		"File containing kubernetes client authentication token. Leave empty to not use a token.")
	flag.StringVar(&clientCertFile, "client-certfile", defaultClientCertFile,
		"File containing client certificate. Leave empty to not use a client certificate.")
	flag.StringVar(&clientKeyFile, "client-keyfile", defaultClientKeyFile,
		"File containing client key. Leave empty to not use a client certificate.")
	flag.IntVar(&ingressPort, "ingress-port", defaultIngressPort,
		"Port to serve ingress traffic to backend services.")
	flag.IntVar(&ingressHealthPort, "ingress-health-port", defaultIngressHealthPort,
		"Port for ingress /health and /status pages. Should be used by frontends to determine if ingress is available.")
	flag.StringVar(&ingressAllow, "ingress-allow", defaultIngressAllow,
		"Source IP or CIDR to allow ingress access by default. This is overridden by the sky.uk/allow "+
			"annotation on ingress resources. Leave empty to deny all access by default.")
	flag.IntVar(&healthPort, "health-port", defaultHealthPort,
		"Port for checking the health of the ingress controller on /health. Also provides /debug/pprof.")
	flag.StringVar(&nginxBinary, "nginx-binary", defaultNginxBinary,
		"Location of nginx binary.")
	flag.StringVar(&nginxWorkDir, "nginx-workdir", defaultNginxWorkingDir,
		"Directory to store nginx files. Also the location of the nginx.tmpl file.")
	flag.IntVar(&nginxWorkerProcesses, "nginx-workers", defaultNginxWorkers,
		"Number of nginx worker processes.")
	flag.IntVar(&nginxWorkerConnections, "nginx-worker-connections", defaultNginxWorkerConnections,
		"Max number of connections per nginx worker. Includes both client and proxy connections.")
	flag.IntVar(&nginxKeepAliveSeconds, "nginx-keepalive-seconds", defaultNginxKeepAliveSeconds,
		"Keep alive time for persistent client connections to nginx. Should generally be set larger than frontend "+
			"keep alive times to prevent stale connections.")
	flag.IntVar(&nginxBackendKeepalives, "nginx-backend-keepalive-count", defaultNginxBackendKeepalives,
		"Maximum number of keepalive connections per backend service. Keepalive connections count against"+
			" nginx-worker-connections limit, and will be restricted by that global limit as well.")
	flag.IntVar(&nginxBackendKeepaliveSeconds, "nginx-backend-keepalive-seconds", defaultNginxBackendKeepaliveSeconds,
		"Time to keep backend keepalive connections open. This should generally be set smaller than backend service keepalive "+
			"times to prevent stale connections.")
	flag.StringVar(&nginxLogLevel, "nginx-loglevel", defaultNginxLogLevel,
		"Log level for nginx. See http://nginx.org/en/docs/ngx_core_module.html#error_log for levels.")
	flag.IntVar(&nginxServerNamesHashBucketSize, "nginx-server-names-hash-bucket-size", defaultNginxServerNamesHashBucketSize,
		"Sets the bucket size for the server names hash tables. Setting this to 0 or less will exclude this "+
			"config from the nginx conf file. The details of setting up hash tables are provided "+
			"in a separate document. http://nginx.org/en/docs/hash.html")
	flag.IntVar(&nginxServerNamesHashMaxSize, "nginx-server-names-hash-max-size", defaultNginxServerNamesHashMaxSize,
		"Sets the maximum size of the server names hash tables. Setting this to 0 or less will exclude this "+
			"config from the nginx conf file. The details of setting up hash tables are provided "+
			"in a separate document. http://nginx.org/en/docs/hash.html")
	flag.StringVar(&nginxTrustedFrontends, "nginx-trusted-frontends", "",
		"Comma separated list of CIDRs to trust when determining the client's real IP from the "+
			"X-Forwarded-For header. The client IP is used for allowing or denying ingress access. "+
			"This will typically be the ELB subnet.")
	flag.StringVar(&elbLabelValue, "elb-label-value", defaultElbLabelValue,
		"Attach to ELBs tagged with "+elb.ElbTag+"=value. Leave empty to not attach.")
	flag.IntVar(&elbExpectedNumber, "elb-expected-number", defaultElbExpectedNumber,
		"Expected number of ELBs to attach to. If 0 the controller will not check,"+
			" otherwise it fails to start if it can't attach to this number.")
	flag.StringVar(&elbRegion, "elb-region", defaultElbRegion,
		"AWS region for ELBs.")
	flag.StringVar(&pushgatewayURL, "pushgateway", "",
		"Prometheus pushgateway URL for pushing metrics. Leave blank to not push metrics.")
	flag.IntVar(&pushgatewayIntervalSeconds, "pushgateway-interval", defaultPushgatewayIntervalSeconds,
		"Interval in seconds for pushing metrics.")
	flag.Var(&pushgatewayLabels, "pushgateway-label",
		"A label=value pair to attach to metrics pushed to prometheus. Specify multiple times for multiple labels.")
}

func main() {
	flag.Parse()

	cmd.ConfigureLogging(debug)
	cmd.ConfigureMetrics("feed-ingress", pushgatewayLabels, pushgatewayURL, pushgatewayIntervalSeconds)

	client := cmd.CreateK8sClient(caCertFile, tokenFile, apiserverURL, clientCertFile, clientKeyFile)
	updaters := createIngressUpdaters()

	controller := controller.New(controller.Config{
		KubernetesClient: client,
		Updaters:         updaters,
		DefaultAllow:     ingressAllow,
	})

	cmd.AddHealthMetrics(controller, metrics.PrometheusIngressSubsystem)
	cmd.AddHealthPort(controller, healthPort)
	cmd.AddSignalHandler(controller)

	err := controller.Start()
	if err != nil {
		log.Error("Error while starting controller: ", err)
		os.Exit(-1)
	}
	log.Info("Controller started")

	select {}
}

func createIngressUpdaters() []controller.Updater {
	frontend := elb.New(elbRegion, elbLabelValue, elbExpectedNumber)
	trustedFrontends := []string{}
	if nginxTrustedFrontends != "" {
		trustedFrontends = strings.Split(nginxTrustedFrontends, ",")
	}
	proxy := nginx.New(nginx.Conf{
		BinaryLocation:            nginxBinary,
		IngressPort:               ingressPort,
		WorkingDir:                nginxWorkDir,
		WorkerProcesses:           nginxWorkerProcesses,
		WorkerConnections:         nginxWorkerConnections,
		KeepaliveSeconds:          nginxKeepAliveSeconds,
		BackendKeepalives:         nginxBackendKeepalives,
		BackendKeepaliveSeconds:   nginxBackendKeepaliveSeconds,
		LogLevel:                  nginxLogLevel,
		ServerNamesHashBucketSize: nginxServerNamesHashBucketSize,
		ServerNamesHashMaxSize:    nginxServerNamesHashMaxSize,
		HealthPort:                ingressHealthPort,
		TrustedFrontends:          trustedFrontends,
	})
	return []controller.Updater{frontend, proxy}
}
<|endoftext|>


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

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"hash/fnv"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"syscall"
	"time"

	consul "github.com/hashicorp/consul/api"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/route"
	"golang.org/x/net/context"

	"github.com/prometheus/prometheus/frankenstein"
	"github.com/prometheus/prometheus/frankenstein/api"
	"github.com/prometheus/prometheus/storage/local"
)

const (
	distributor = "distributor"
	ingestor    = "ingestor"
	infName     = "eth0"
)

func main() {
	var (
		mode                 string
		listenPort           int
		consulHost           string
		consulPrefix         string
		s3URL                string
		dynamodbURL          string
		dynamodbCreateTables bool
		memcachedHostname    string
		memcachedTimeout     time.Duration
		memcachedExpiration  time.Duration
		memcachedService     string
		remoteTimeout        time.Duration
	)

	flag.StringVar(&mode, "mode", distributor, "Mode (distributor, ingestor).")
	flag.IntVar(&listenPort, "web.listen-port", 9094, "HTTP server listen port.")
	flag.StringVar(&consulHost, "consul.hostname", "localhost:8500", "Hostname and port of Consul.")
	flag.StringVar(&consulPrefix, "consul.prefix", "collectors/", "Prefix for keys in Consul.")
	flag.StringVar(&s3URL, "s3.url", "localhost:4569", "S3 endpoint URL.")
	flag.StringVar(&dynamodbURL, "dynamodb.url", "localhost:8000", "DynamoDB endpoint URL.")
	flag.BoolVar(&dynamodbCreateTables, "dynamodb.create-tables", false, "Create required DynamoDB tables on startup.")
	flag.StringVar(&memcachedHostname, "memcached.hostname", "", "Hostname for memcached service to use when caching chunks. If empty, no memcached will be used.")
	flag.DurationVar(&memcachedTimeout, "memcached.timeout", 100*time.Millisecond, "Maximum time to wait before giving up on memcached requests.")
	flag.DurationVar(&memcachedExpiration, "memcached.expiration", 2*15*time.Second, "How long chunks stay in the memcache.")
	flag.StringVar(&memcachedService, "memcached.service", "memcached", "SRV service used to discover memcache servers.")
	flag.DurationVar(&remoteTimeout, "remote.timeout", 100*time.Millisecond, "Timeout for downstream ingestors.")
	flag.Parse()

	consul, err := frankenstein.NewConsulClient(consulHost)
	if err != nil {
		log.Fatalf("Error initializing Consul client: %v", err)
	}

	var memcache *frankenstein.MemcacheClient
	if memcachedHostname != "" {
		memcache = frankenstein.NewMemcacheClient(frankenstein.MemcacheConfig{
			Host:           memcachedHostname,
			Service:        memcachedService,
			Timeout:        memcachedTimeout,
			UpdateInterval: 1 * time.Minute,
		})
	}
	chunkStore, err := frankenstein.NewAWSChunkStore(frankenstein.ChunkStoreConfig{
		S3URL:          s3URL,
		DynamoDBURL:    dynamodbURL,
		MemcacheClient: memcache,
	})
	if err != nil {
		log.Fatal(err)
	}
	if dynamodbCreateTables {
		if err = chunkStore.CreateTables(); err != nil {
			log.Fatal(err)
		}
	}

	switch mode {
	case distributor:
		setupDistributor(consul, consulPrefix, chunkStore, remoteTimeout)
	case ingestor:
		ingestor := setupIngestor(consul, consulPrefix, chunkStore, listenPort)
		defer ingestor.Stop()
		defer deleteIngestorConfigFromConsul(consul, consulPrefix)
	default:
		log.Fatalf("Mode %s not supported!", mode)
	}

	http.Handle("/metrics", prometheus.Handler())
	go http.ListenAndServe(fmt.Sprintf(":%d", listenPort), nil)

	term := make(chan os.Signal)
	signal.Notify(term, os.Interrupt, syscall.SIGTERM)
	<-term
	log.Warn("Received SIGTERM, exiting gracefully...")
}

func setupDistributor(
	consulClient frankenstein.ConsulClient,
	consulPrefix string,
	chunkStore frankenstein.ChunkStore,
	remoteTimeout time.Duration,
) {
	clientFactory := func(hostname string) (*frankenstein.IngesterClient, error) {
		return frankenstein.NewIngesterClient(hostname, remoteTimeout), nil
	}
	distributor, err := frankenstein.NewDistributor(frankenstein.DistributorConfig{
		ConsulClient:  consulClient,
		ConsulPrefix:  consulPrefix,
		ClientFactory: clientFactory,
	})
	if err != nil {
		log.Fatal(err)
	}
	http.Handle("/api/prom/push", frankenstein.AppenderHandler(distributor))

	// This sets up a complete querying pipeline:
	//
	// PromQL -> MergeQuerier -> Distributor -> IngesterQuerier -> Ingester
	//              |
	//              +----------> ChunkQuerier -> DynamoDB/S3
	//
	// TODO: Move querier to separate binary.
	newQuerier := func(ctx context.Context) local.Querier {
		return frankenstein.MergeQuerier{
			Queriers: []frankenstein.Querier{
				distributor,
				&frankenstein.ChunkQuerier{
					Store: chunkStore,
				},
			},
			Context: ctx,
		}
	}

	api := api.New(newQuerier)
	router := route.New()
	api.Register(router.WithPrefix("/api/prom/api/v1"))
	http.Handle("/", router)
}

func setupIngestor(
	consulClient frankenstein.ConsulClient,
	consulPrefix string,
	chunkStore frankenstein.ChunkStore,
	listenPort int,
) *local.Ingestor {
	for i := 0; i < 10; i++ {
		if err := writeIngestorConfigToConsul(consulClient, consulPrefix, listenPort); err == nil {
			break
		} else {
			log.Errorf("Failed to write to consul, sleeping: %v", err)
			time.Sleep(1 * time.Second)
		}
	}

	ingestor, err := local.NewIngestor(chunkStore)
	if err != nil {
		log.Fatal(err)
	}
	prometheus.MustRegister(ingestor)
	http.Handle("/push", frankenstein.AppenderHandler(ingestor))
	http.Handle("/query", frankenstein.QueryHandler(ingestor))
	http.Handle("/label_values", frankenstein.LabelValuesHandler(ingestor))
	return ingestor
}

// GetFirstAddressOf returns the first IPv4 address of the supplied interface name.
func GetFirstAddressOf(name string) (string, error) {
	inf, err := net.InterfaceByName(name)
	if err != nil {
		return "", err
	}

	addrs, err := inf.Addrs()
	if err != nil {
		return "", err
	}
	if len(addrs) <= 0 {
		return "", fmt.Errorf("No address found for %s", name)
	}

	for _, addr := range addrs {
		switch v := addr.(type) {
		case *net.IPNet:
			if ip := v.IP.To4(); ip != nil {
				return v.IP.String(), nil
			}
		}
	}

	return "", fmt.Errorf("No address found for %s", name)
}

func writeIngestorConfigToConsul(consulClient frankenstein.ConsulClient, consulPrefix string, listenPort int) error {
	log.Info("Adding ingestor to consul")

	hostname, err := os.Hostname()
	if err != nil {
		return err
	}

	addr, err := GetFirstAddressOf(infName)
	if err != nil {
		return err
	}

	tokenHasher := fnv.New64()
	tokenHasher.Write([]byte(hostname))

	buf, err := json.Marshal(frankenstein.Collector{
		ID:       hostname,
		Hostname: fmt.Sprintf("%s:%d", addr, listenPort),
		Tokens:   []uint64{tokenHasher.Sum64()},
	})
	if err != nil {
		return err
	}

	_, err = consulClient.Put(&consul.KVPair{
		Key:   consulPrefix + hostname,
		Value: buf,
	}, &consul.WriteOptions{})
	return err
}

func deleteIngestorConfigFromConsul(consulClient frankenstein.ConsulClient, consulPrefix string) error {
	log.Info("Removing ingestor from consul")

	hostname, err := os.Hostname()
	if err != nil {
		return err
	}

	buf, err := json.Marshal(frankenstein.Collector{
		ID:       hostname,
		Hostname: "",
		Tokens:   []uint64{},
	})
	if err != nil {
		return err
	}

	_, err = consulClient.Put(&consul.KVPair{
		Key:   consulPrefix + hostname,
		Value: buf,
	}, &consul.WriteOptions{})
	return err
}
<commit_msg>Increase distributor->ingestor timeout. (#47)<commit_after>// Copyright 2016 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"hash/fnv"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"syscall"
	"time"

	consul "github.com/hashicorp/consul/api"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/route"
	"golang.org/x/net/context"

	"github.com/prometheus/prometheus/frankenstein"
	"github.com/prometheus/prometheus/frankenstein/api"
	"github.com/prometheus/prometheus/storage/local"
)

const (
	distributor = "distributor"
	ingestor    = "ingestor"
	infName     = "eth0"
)

func main() {
	var (
		mode                 string
		listenPort           int
		consulHost           string
		consulPrefix         string
		s3URL                string
		dynamodbURL          string
		dynamodbCreateTables bool
		memcachedHostname    string
		memcachedTimeout     time.Duration
		memcachedExpiration  time.Duration
		memcachedService     string
		remoteTimeout        time.Duration
		flushPeriod          time.Duration
		maxChunkAge          time.Duration
	)

	flag.StringVar(&mode, "mode", distributor, "Mode (distributor, ingestor).")
	flag.IntVar(&listenPort, "web.listen-port", 9094, "HTTP server listen port.")
	flag.StringVar(&consulHost, "consul.hostname", "localhost:8500", "Hostname and port of Consul.")
	flag.StringVar(&consulPrefix, "consul.prefix", "collectors/", "Prefix for keys in Consul.")
	flag.StringVar(&s3URL, "s3.url", "localhost:4569", "S3 endpoint URL.")
	flag.StringVar(&dynamodbURL, "dynamodb.url", "localhost:8000", "DynamoDB endpoint URL.")
	flag.BoolVar(&dynamodbCreateTables, "dynamodb.create-tables", false, "Create required DynamoDB tables on startup.")
	flag.StringVar(&memcachedHostname, "memcached.hostname", "", "Hostname for memcached service to use when caching chunks. If empty, no memcached will be used.")
	flag.DurationVar(&memcachedTimeout, "memcached.timeout", 100*time.Millisecond, "Maximum time to wait before giving up on memcached requests.")
	flag.DurationVar(&memcachedExpiration, "memcached.expiration", 2*15*time.Second, "How long chunks stay in the memcache.")
	flag.StringVar(&memcachedService, "memcached.service", "memcached", "SRV service used to discover memcache servers.")
	flag.DurationVar(&remoteTimeout, "remote.timeout", 5*time.Second, "Timeout for downstream ingestors.")
	flag.DurationVar(&flushPeriod, "ingestor.flush-period", 1*time.Minute, "Period with which to attempt to flush chunks.")
	flag.DurationVar(&maxChunkAge, "ingestor.max-chunk-age", 10*time.Minute, "Maximum chunk age before flushing.")
	flag.Parse()

	consul, err := frankenstein.NewConsulClient(consulHost)
	if err != nil {
		log.Fatalf("Error initializing Consul client: %v", err)
	}

	var memcache *frankenstein.MemcacheClient
	if memcachedHostname != "" {
		memcache = frankenstein.NewMemcacheClient(frankenstein.MemcacheConfig{
			Host:           memcachedHostname,
			Service:        memcachedService,
			Timeout:        memcachedTimeout,
			UpdateInterval: 1 * time.Minute,
		})
	}
	chunkStore, err := frankenstein.NewAWSChunkStore(frankenstein.ChunkStoreConfig{
		S3URL:          s3URL,
		DynamoDBURL:    dynamodbURL,
		MemcacheClient: memcache,
	})
	if err != nil {
		log.Fatal(err)
	}
	if dynamodbCreateTables {
		if err = chunkStore.CreateTables(); err != nil {
			log.Fatal(err)
		}
	}

	switch mode {
	case distributor:
		setupDistributor(consul, consulPrefix, chunkStore, remoteTimeout)
	case ingestor:
		cfg := local.IngestorConfig{
			FlushCheckPeriod: flushPeriod,
			MaxChunkAge:      maxChunkAge,
		}
		ingestor := setupIngestor(consul, consulPrefix, chunkStore, listenPort, cfg)
		defer ingestor.Stop()
		defer deleteIngestorConfigFromConsul(consul, consulPrefix)
	default:
		log.Fatalf("Mode %s not supported!", mode)
	}

	http.Handle("/metrics", prometheus.Handler())
	go http.ListenAndServe(fmt.Sprintf(":%d", listenPort), nil)

	term := make(chan os.Signal)
	signal.Notify(term, os.Interrupt, syscall.SIGTERM)
	<-term
	log.Warn("Received SIGTERM, exiting gracefully...")
}

func setupDistributor(
	consulClient frankenstein.ConsulClient,
	consulPrefix string,
	chunkStore frankenstein.ChunkStore,
	remoteTimeout time.Duration,
) {
	clientFactory := func(hostname string) (*frankenstein.IngesterClient, error) {
		return frankenstein.NewIngesterClient(hostname, remoteTimeout), nil
	}
	distributor, err := frankenstein.NewDistributor(frankenstein.DistributorConfig{
		ConsulClient:  consulClient,
		ConsulPrefix:  consulPrefix,
		ClientFactory: clientFactory,
	})
	if err != nil {
		log.Fatal(err)
	}
	http.Handle("/api/prom/push", frankenstein.AppenderHandler(distributor))

	// This sets up a complete querying pipeline:
	//
	// PromQL -> MergeQuerier -> Distributor -> IngesterQuerier -> Ingester
	//              |
	//              +----------> ChunkQuerier -> DynamoDB/S3
	//
	// TODO: Move querier to separate binary.
	newQuerier := func(ctx context.Context) local.Querier {
		return frankenstein.MergeQuerier{
			Queriers: []frankenstein.Querier{
				distributor,
				&frankenstein.ChunkQuerier{
					Store: chunkStore,
				},
			},
			Context: ctx,
		}
	}

	api := api.New(newQuerier)
	router := route.New()
	api.Register(router.WithPrefix("/api/prom/api/v1"))
	http.Handle("/", router)
}

func setupIngestor(
	consulClient frankenstein.ConsulClient,
	consulPrefix string,
	chunkStore frankenstein.ChunkStore,
	listenPort int,
	cfg local.IngestorConfig,
) *local.Ingestor {
	for i := 0; i < 10; i++ {
		if err := writeIngestorConfigToConsul(consulClient, consulPrefix, listenPort); err == nil {
			break
		} else {
			log.Errorf("Failed to write to consul, sleeping: %v", err)
			time.Sleep(1 * time.Second)
		}
	}

	ingestor, err := local.NewIngestor(cfg, chunkStore)
	if err != nil {
		log.Fatal(err)
	}
	prometheus.MustRegister(ingestor)
	http.Handle("/push", frankenstein.AppenderHandler(ingestor))
	http.Handle("/query", frankenstein.QueryHandler(ingestor))
	http.Handle("/label_values", frankenstein.LabelValuesHandler(ingestor))
	return ingestor
}

// GetFirstAddressOf returns the first IPv4 address of the supplied interface name.
func GetFirstAddressOf(name string) (string, error) {
	inf, err := net.InterfaceByName(name)
	if err != nil {
		return "", err
	}

	addrs, err := inf.Addrs()
	if err != nil {
		return "", err
	}
	if len(addrs) <= 0 {
		return "", fmt.Errorf("No address found for %s", name)
	}

	for _, addr := range addrs {
		switch v := addr.(type) {
		case *net.IPNet:
			if ip := v.IP.To4(); ip != nil {
				return v.IP.String(), nil
			}
		}
	}

	return "", fmt.Errorf("No address found for %s", name)
}

func writeIngestorConfigToConsul(consulClient frankenstein.ConsulClient, consulPrefix string, listenPort int) error {
	log.Info("Adding ingestor to consul")

	hostname, err := os.Hostname()
	if err != nil {
		return err
	}

	addr, err := GetFirstAddressOf(infName)
	if err != nil {
		return err
	}

	tokenHasher := fnv.New64()
	tokenHasher.Write([]byte(hostname))

	buf, err := json.Marshal(frankenstein.Collector{
		ID:       hostname,
		Hostname: fmt.Sprintf("%s:%d", addr, listenPort),
		Tokens:   []uint64{tokenHasher.Sum64()},
	})
	if err != nil {
		return err
	}

	_, err = consulClient.Put(&consul.KVPair{
		Key:   consulPrefix + hostname,
		Value: buf,
	}, &consul.WriteOptions{})
	return err
}

func deleteIngestorConfigFromConsul(consulClient frankenstein.ConsulClient, consulPrefix string) error {
	log.Info("Removing ingestor from consul")

	hostname, err := os.Hostname()
	if err != nil {
		return err
	}

	buf, err := json.Marshal(frankenstein.Collector{
		ID:       hostname,
		Hostname: "",
		Tokens:   []uint64{},
	})
	if err != nil {
		return err
	}

	_, err = consulClient.Put(&consul.KVPair{
		Key:   consulPrefix + hostname,
		Value: buf,
	}, &consul.WriteOptions{})
	return err
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"sync/atomic"
	"time"

	"github.com/gocolly/colly"
	"github.com/gocolly/colly/extensions"
)

type Config struct {
	WaitSeconds      int
	LimitFlag        int
	GameList         *FileManager
	Errors           ErrorList
	totalPages       uint32
	MainCollector    *colly.Collector
	InitialCollector *colly.Collector
}

func (c *Config) GetTotalPages() int {

	return int(atomic.LoadUint32(&c.totalPages))
}

func (c *Config) SetTotalPages(pages int) {
	atomic.StoreUint32(&c.totalPages, uint32(pages))
}

func (c *Config) GeekURL(num int) string {
	return fmt.Sprintf("https://boardgamegeek.com/browse/boardgame/page/%v", num)
}

func NewConfig() *Config {
	os.Mkdir("data", 0777)
	var (
		waitSeconds int
		limitFlag   int
	)
	flag.IntVar(&waitSeconds, "wait", 5, "wait sets the default wait between web requests")
	flag.IntVar(&limitFlag, "limit", 0, "limit sets the max number of pages to check (0 is no limit)")
	flag.Parse()
	config := &Config{
		WaitSeconds: waitSeconds,
		LimitFlag:   limitFlag,
		GameList:    &FileManager{filename: filepath.Join("data", "games.txt")},
		Errors:      ErrorList{},
		totalPages:  0,
		MainCollector: colly.NewCollector(
			colly.Async(waitSeconds == 0),
		),
		InitialCollector: colly.NewCollector(),
	}

	extensions.RandomUserAgent(config.InitialCollector)
	extensions.RandomUserAgent(config.MainCollector)
	if waitSeconds > 0 {
		config.MainCollector.Limit(&colly.LimitRule{
			Delay: time.Duration(waitSeconds) * time.Second,
		})
	} else {
		config.MainCollector.Limit(&colly.LimitRule{
			Delay:       time.Duration(0) * time.Second,
			Parallelism: 20,
			DomainGlob:  "*",
		})
	}

	// set error handling
	re := regexp.MustCompile("page/([0-9]*)(?:|$)")
	responseHandler := func(r *colly.Response) {
		if r.StatusCode != 200 {
			pageString := re.FindStringSubmatch(r.Request.URL.String())[1]
			page, err := strconv.Atoi(pageString)
			if err != nil {
				log.Println("parsing page string", err)
			}
			e := HttpResponseError{
				Code: r.StatusCode,
				Body: string(r.Body),
				Page: page,
			}
			log.Println("Error", e.Code, e.Body)
			config.Errors.Add(e)
		}
	}
	config.InitialCollector.OnResponse(responseHandler)
	config.MainCollector.OnResponse(responseHandler)
	errHandler := func(r *colly.Response, responseErr error) {
		pageString := re.FindStringSubmatch(r.Request.URL.String())[1]
		page, err := strconv.Atoi(pageString)
		if err != nil {
			fmt.Println("error parsing '", pageString, "' from url", r.Request.URL.String())
			return
		}
		e := HttpResponseError{
			Code:  r.StatusCode,
			Body:  string(r.Body),
			Error: responseErr.Error(),
			Page:  page,
		}
		log.Println("Error", e.Code, e.Error)
		config.Errors.Add(e)
	}
	config.MainCollector.OnError(errHandler)
	config.InitialCollector.OnError(errHandler)

	return config
}
<commit_msg>allow more simultaneous requests to same host<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"sync/atomic"
	"time"

	"github.com/gocolly/colly"
	"github.com/gocolly/colly/extensions"
)

type Config struct {
	WaitSeconds      int
	LimitFlag        int
	GameList         *FileManager
	Errors           ErrorList
	totalPages       uint32
	MainCollector    *colly.Collector
	InitialCollector *colly.Collector
}

func (c *Config) GetTotalPages() int {

	return int(atomic.LoadUint32(&c.totalPages))
}

func (c *Config) SetTotalPages(pages int) {
	atomic.StoreUint32(&c.totalPages, uint32(pages))
}

func (c *Config) GeekURL(num int) string {
	return fmt.Sprintf("https://boardgamegeek.com/browse/boardgame/page/%v", num)
}

func NewConfig() *Config {
	os.Mkdir("data", 0777)
	var (
		waitSeconds int
		limitFlag   int
	)
	flag.IntVar(&waitSeconds, "wait", 5, "wait sets the default wait between web requests")
	flag.IntVar(&limitFlag, "limit", 0, "limit sets the max number of pages to check (0 is no limit)")
	flag.Parse()
	config := &Config{
		WaitSeconds: waitSeconds,
		LimitFlag:   limitFlag,
		GameList:    &FileManager{filename: filepath.Join("data", "games.txt")},
		Errors:      ErrorList{},
		totalPages:  0,
		MainCollector: colly.NewCollector(
			colly.Async(waitSeconds == 0),
		),
		InitialCollector: colly.NewCollector(),
	}

	extensions.RandomUserAgent(config.InitialCollector)
	extensions.RandomUserAgent(config.MainCollector)
	if waitSeconds > 0 {
		config.MainCollector.Limit(&colly.LimitRule{
			Delay: time.Duration(waitSeconds) * time.Second,
		})
	} else {
		config.MainCollector.Limit(&colly.LimitRule{
			Delay:       time.Duration(0) * time.Second,
			Parallelism: 20,
			DomainGlob:  "*",
		})
		// Some hints on transport tweaks for scraping
		// http://tleyden.github.io/blog/2016/11/21/tuning-the-go-http-client-library-for-load-testing/
		defaultRoundTripper := http.DefaultTransport
		defaultRoundTripperPointer := defaultRoundTripper.(*http.Transport)
		t := *defaultRoundTripperPointer // deref to get copy
		t.MaxIdleConns = 0
		t.MaxIdleConnsPerHost = 1000
		config.MainCollector.WithTransport(&t)
	}

	// set error handling
	re := regexp.MustCompile("page/([0-9]*)(?:|$)")
	responseHandler := func(r *colly.Response) {
		if r.StatusCode != 200 {
			pageString := re.FindStringSubmatch(r.Request.URL.String())[1]
			page, err := strconv.Atoi(pageString)
			if err != nil {
				log.Println("parsing page string", err)
			}
			e := HttpResponseError{
				Code: r.StatusCode,
				Body: string(r.Body),
				Page: page,
			}
			log.Println("Error", e.Code, e.Body)
			config.Errors.Add(e)
		}
	}
	config.InitialCollector.OnResponse(responseHandler)
	config.MainCollector.OnResponse(responseHandler)
	errHandler := func(r *colly.Response, responseErr error) {
		pageString := re.FindStringSubmatch(r.Request.URL.String())[1]
		page, err := strconv.Atoi(pageString)
		if err != nil {
			fmt.Println("error parsing '", pageString, "' from url", r.Request.URL.String())
			return
		}
		e := HttpResponseError{
			Code:  r.StatusCode,
			Body:  string(r.Body),
			Error: responseErr.Error(),
			Page:  page,
		}
		log.Println("Error", e.Code, e.Error)
		config.Errors.Add(e)
	}
	config.MainCollector.OnError(errHandler)
	config.InitialCollector.OnError(errHandler)

	return config
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"sort"
	"strings"

	"github.com/sprucehealth/graphql/language/ast"
)

func generateClient(g *generator) {
	imports := []string{"github.com/sprucehealth/graphql"}
	if len(g.cfg.Resolvers) != 0 {
		imports = []string{
			"context",
			"fmt",
			"reflect",
			"",
			"github.com/sprucehealth/graphql",
			"github.com/sprucehealth/graphql/gqldecode",
			"github.com/sprucehealth/graphql/gqlerrors",
			"github.com/sprucehealth/graphql/language/parser",
			"github.com/sprucehealth/graphql/language/printer",
			"github.com/sprucehealth/graphql/language/source",
			"github.com/sprucehealth/mapstructure",
		}
	}

	g.printf("package client\n\n")
	g.printf("import (\n")
	for _, im := range imports {
		if im == "" {
			g.printf("\n")
		} else {
			g.printf("\t%q\n", im)
		}
	}
	g.printf(")\n\n")

	clientTypes := make(map[string]struct{})
	for _, t := range strings.Split(*flagClientTypes, ",") {
		clientTypes[t] = struct{}{}
	}

	var clientTypeDefs []*ast.ObjectDefinition
	for _, def := range g.doc.Definitions {
		switch def := def.(type) {
		case *ast.ObjectDefinition:
			if _, ok := clientTypes[def.Name.Value]; ok {
				clientTypeDefs = append(clientTypeDefs, def)
			}
			g.genObjectModel(def)
			g.printf("\n")
		case *ast.InterfaceDefinition:
			g.genInterfaceModel(def)
			g.printf("\n")
		case *ast.InputObjectDefinition:
			g.genInputModel(def)
			g.printf("\n")
		case *ast.UnionDefinition:
			g.genUnionModel(def)
			g.printf("\n")
		case *ast.EnumDefinition:
			g.genEnumConstants(def)
			g.printf("\n")
		}
	}

	var sigs []string
	for _, def := range clientTypeDefs {
		for _, field := range def.Fields {
			sigs = append(sigs, renderSignatureForField(g, def, field))
		}
	}
	sort.Strings(sigs)
	g.printf("type Client interface {\n")
	for _, sig := range sigs {
		g.printf(sig + "\n")
	}
	g.printf("}\n")
	g.print(`
		type client struct {
			endpoint string
			authToken string
			headers map[string][]string
			log Logger
		}

		type clientOption func(c *client)

		func WithClientLogger(l Logger) clientOption {
			return func(c *client) {
				c.log = l
			}
		}

		func WithHeaders(h map[string][]string) clientOption {
			return func(c *client) {
				c.headers = h
			}
		}

		type Logger interface {
			Debugf(ctx context.Context, msg string, v ...interface{})
		}

		type nullLogger struct{}
		func (nullLogger) Debugf(ctx context.Context, msg string, v ...interface{}){}

		func New(endpoint, authToken string, opts ...clientOption) Client {
			c := &client{
				endpoint: endpoint,
				authToken: authToken,
				headers: make(map[string][]string),
				log: nullLogger{},
			}
			for _, o := range opts {
				o(c)
			}
			return c
		}
	`)
	for _, def := range clientTypeDefs {
		for _, field := range def.Fields {
			genClientMethodForField(g, def, field)
			g.printf("\n")
		}
	}
	genDecoderTypeMap(g)
	g.printf("\n")
	genDecoderHook(g)
	g.printf("\n")
	genQueryWrapperTypes(g)
	g.printf("\n")
	genClientDo(g)
	g.printf("\n")
	genRewriteQuery(g)
	g.printf("\n")
}

func renderSignatureForField(g *generator, d *ast.ObjectDefinition, f *ast.FieldDefinition) string {
	return fmt.Sprintf("%s%s(ctx context.Context, query string) (%s, error)", exportedName(d.Name.Value), exportedName(f.Name.Value), g.goType(f.Type, ""))
}

func genDecoderHook(g *generator) {
	g.printf("func decoderHook(from reflect.Kind, to reflect.Kind, v interface{}) (interface{}, error) {\n")
	g.printf("\tif from == reflect.Map && to == reflect.Interface {\n")
	g.printf("\t\tcVal := reflect.New(objectTypesByTypename[v.(map[string]interface{})[\"__typename\"].(string)])\n")
	g.printf("\t\tout := cVal.Elem().Addr().Interface()\n")
	g.printf("\t\tdecoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{\n")
	g.printf("\t\t\tDecodeHook: decoderHook,\n")
	g.printf("\t\t\tResult:     out,\n")
	g.printf("\t\t})\n")
	g.printf("\t\tif err != nil {\n")
	g.printf("\t\t\treturn 0, err\n")
	g.printf("\t\t}\n")
	g.printf("\t\tif err := decoder.Decode(v.(map[string]interface{})); err != nil {\n")
	g.print("\t\t\treturn nil, fmt.Errorf(\"Error decoding %+v into %+v: %s\", v, out, err)\n")
	g.printf("\t\t}\n")
	g.printf("\t\treturn out, nil\n")
	g.printf("\t}\n")
	g.printf("\treturn v, nil\n")
	g.printf("}\n")
}

func genDecoderTypeMap(g *generator) {
	g.printf("var objectTypesByTypename = map[string]reflect.Type{\n")
	for _, def := range g.doc.Definitions {
		switch def := def.(type) {
		case *ast.ObjectDefinition:
			g.printf("\"%s\": reflect.TypeOf(%s{}),\n", def.Name.Value, def.Name.Value)
		}
	}
	g.printf("}\n")
}

func genOutputTypeVar(g *generator, f *ast.FieldDefinition) {
	outType := g.goType(f.Type, "")
	if outType[0] == '*' {
		outType = outType[1:]
	}
	g.printf("\tvar out %s\n", outType)
}

func outputTypeReturn(g *generator, f *ast.FieldDefinition) string {
	outType := g.goType(f.Type, "")
	if outType[0] == '*' {
		return "&out"
	}
	return "out"
}

func genClientMethodForField(g *generator, d *ast.ObjectDefinition, f *ast.FieldDefinition) {
	g.printf("func (c *client) %s {\n", renderSignatureForField(g, d, f))
	g.printf("\tc.log.Debugf(ctx, \"%s%s\")\n", exportedName(d.Name.Value), exportedName(f.Name.Value))
	genOutputTypeVar(g, f)
	g.printf("\tif _, err := c.do(ctx, \"%s\", query, &out); err != nil {\n", unexportedName(f.Name.Value))
	g.printf("\t\treturn %s, err\n", renderDefaultReturnValue(f.Type))
	g.printf("\t}\n")
	g.printf("\treturn %s, nil", outputTypeReturn(g, f))
	g.printf("}\n")
}

func renderDefaultClientReturnValue(t ast.Type) string {
	switch t := t.(type) {
	case *ast.NonNull:
		return renderDefaultClientReturnValue(t.Type)
	case *ast.Named:
		switch t.Name.Value {
		case "ID", "String":
			return `""`
		case "Boolean":
			return "false"
		case "Float", "Int":
			return "0"
		}
	}
	return "nil"
}

func genQueryWrapperTypes(g *generator) {
	g.print(`
		type gqlRequestBody struct {
			Query string ` + "`json:\"query\"`" + `
		}

		type gqlResponse struct {
			Data   map[string]interface{}   ` + "`json:\"data\"`" + `
			Errors []map[string]interface{} ` + "`json:\"errors\"`" + `
		}
	`)
}

func genRewriteQuery(g *generator) {
	g.print(`
		func rewriteQuery(query string) (string, error) {
			qast, err := parser.Parse(parser.ParseParams{
				Source: source.New("GraphQL Query", query),
			})
			if err != nil {
				return "", err
			}
			graphql.RequestTypeNames(qast)
			return printer.Print(qast), nil
		}
	`)
}

func genClientDo(g *generator) {
	g.print(`
		func (c *client) do(ctx context.Context, dataField, query string, out interface{}) (int, error) {
			query, err := rewriteQuery(query)
			if err != nil {
				return 0, err
			}
			rb := &gqlRequestBody{Query: query}
			bBody, err := json.Marshal(rb)
			if err != nil {
				return 0, err
			}
			c.log.Debugf(ctx, "Request: %s - %s", c.endpoint, string(bBody))
			req, err := http.NewRequest(http.MethodPost, c.endpoint, bytes.NewReader(bBody))
			if err != nil {
				return 0, err
			}
			if c.authToken != "" {
				req.AddCookie(&http.Cookie{
					Name:     "at",
					Value:    c.authToken,
				})
			}
			for h, hvs := range c.headers {
				for _, hv := range hvs {
					req.Header.Add(h, hv)
				}
			}
			resp, err := http.DefaultClient.Do(req)
			if err != nil {
				return 0, err
			}
			ball, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return resp.StatusCode, fmt.Errorf("Error reading body - in response from %v: %s", req.URL.String(), err)
			}
			c.log.Debugf(ctx, "Response: %s - %s", resp.Status, ball)
			gqlResp := &gqlResponse{}
			if resp.StatusCode == http.StatusOK {
				if err := json.NewDecoder(bytes.NewReader(ball)).Decode(gqlResp); err != nil {
					return resp.StatusCode, fmt.Errorf("Error parsing body into output: %s", err)
				}
				if len(gqlResp.Errors) != 0 {
					var allErrors string
					for i, err := range gqlResp.Errors {
						allErrors += fmt.Sprintf("%d. - %+v\n", i, err)
					}
					return resp.StatusCode, errors.New(allErrors)
				}
				decoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{
					DecodeHook: decoderHook,
					Result:     out,
				})
				if err != nil {
					return 0, err
				}
				mData, ok := gqlResp.Data[dataField].(map[string]interface{})
				if ok {
					if err := decoder.Decode(mData); err != nil {
						return resp.StatusCode, fmt.Errorf("Error parsing body into output: %s", err)
					}
				} else {
					sData, ok := gqlResp.Data[dataField].([]interface{})
					if ok {
						if err := decoder.Decode(sData); err != nil {
							return resp.StatusCode, fmt.Errorf("Error parsing body into output: %s", err)
						}
					} else {
						return resp.StatusCode, fmt.Errorf("unhandled response data type %T %+v", gqlResp.Data[dataField], gqlResp.Data[dataField])
					}
				}

				return resp.StatusCode, nil
			}
			return resp.StatusCode, fmt.Errorf("Non 200 Response (%d) from %s: %s", resp.StatusCode, req.URL, string(ball))
		}
	`)
}
<commit_msg>graphql2go: make sure to close http body in client<commit_after>package main

import (
	"fmt"
	"sort"
	"strings"

	"github.com/sprucehealth/graphql/language/ast"
)

func generateClient(g *generator) {
	imports := []string{"github.com/sprucehealth/graphql"}
	if len(g.cfg.Resolvers) != 0 {
		imports = []string{
			"context",
			"fmt",
			"reflect",
			"",
			"github.com/sprucehealth/graphql",
			"github.com/sprucehealth/graphql/gqldecode",
			"github.com/sprucehealth/graphql/gqlerrors",
			"github.com/sprucehealth/graphql/language/parser",
			"github.com/sprucehealth/graphql/language/printer",
			"github.com/sprucehealth/graphql/language/source",
			"github.com/sprucehealth/mapstructure",
		}
	}

	g.printf("package client\n\n")
	g.printf("import (\n")
	for _, im := range imports {
		if im == "" {
			g.printf("\n")
		} else {
			g.printf("\t%q\n", im)
		}
	}
	g.printf(")\n\n")

	clientTypes := make(map[string]struct{})
	for _, t := range strings.Split(*flagClientTypes, ",") {
		clientTypes[t] = struct{}{}
	}

	var clientTypeDefs []*ast.ObjectDefinition
	for _, def := range g.doc.Definitions {
		switch def := def.(type) {
		case *ast.ObjectDefinition:
			if _, ok := clientTypes[def.Name.Value]; ok {
				clientTypeDefs = append(clientTypeDefs, def)
			}
			g.genObjectModel(def)
			g.printf("\n")
		case *ast.InterfaceDefinition:
			g.genInterfaceModel(def)
			g.printf("\n")
		case *ast.InputObjectDefinition:
			g.genInputModel(def)
			g.printf("\n")
		case *ast.UnionDefinition:
			g.genUnionModel(def)
			g.printf("\n")
		case *ast.EnumDefinition:
			g.genEnumConstants(def)
			g.printf("\n")
		}
	}

	var sigs []string
	for _, def := range clientTypeDefs {
		for _, field := range def.Fields {
			sigs = append(sigs, renderSignatureForField(g, def, field))
		}
	}
	sort.Strings(sigs)
	g.printf("type Client interface {\n")
	for _, sig := range sigs {
		g.printf(sig + "\n")
	}
	g.printf("}\n")
	g.print(`
		type client struct {
			endpoint string
			authToken string
			headers map[string][]string
			log Logger
		}

		type clientOption func(c *client)

		func WithClientLogger(l Logger) clientOption {
			return func(c *client) {
				c.log = l
			}
		}

		func WithHeaders(h map[string][]string) clientOption {
			return func(c *client) {
				c.headers = h
			}
		}

		type Logger interface {
			Debugf(ctx context.Context, msg string, v ...interface{})
		}

		type nullLogger struct{}
		func (nullLogger) Debugf(ctx context.Context, msg string, v ...interface{}){}

		func New(endpoint, authToken string, opts ...clientOption) Client {
			c := &client{
				endpoint: endpoint,
				authToken: authToken,
				headers: make(map[string][]string),
				log: nullLogger{},
			}
			for _, o := range opts {
				o(c)
			}
			return c
		}
	`)
	for _, def := range clientTypeDefs {
		for _, field := range def.Fields {
			genClientMethodForField(g, def, field)
			g.printf("\n")
		}
	}
	genDecoderTypeMap(g)
	g.printf("\n")
	genDecoderHook(g)
	g.printf("\n")
	genQueryWrapperTypes(g)
	g.printf("\n")
	genClientDo(g)
	g.printf("\n")
	genRewriteQuery(g)
	g.printf("\n")
}

func renderSignatureForField(g *generator, d *ast.ObjectDefinition, f *ast.FieldDefinition) string {
	return fmt.Sprintf("%s%s(ctx context.Context, query string) (%s, error)", exportedName(d.Name.Value), exportedName(f.Name.Value), g.goType(f.Type, ""))
}

func genDecoderHook(g *generator) {
	g.printf("func decoderHook(from reflect.Kind, to reflect.Kind, v interface{}) (interface{}, error) {\n")
	g.printf("\tif from == reflect.Map && to == reflect.Interface {\n")
	g.printf("\t\tcVal := reflect.New(objectTypesByTypename[v.(map[string]interface{})[\"__typename\"].(string)])\n")
	g.printf("\t\tout := cVal.Elem().Addr().Interface()\n")
	g.printf("\t\tdecoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{\n")
	g.printf("\t\t\tDecodeHook: decoderHook,\n")
	g.printf("\t\t\tResult:     out,\n")
	g.printf("\t\t})\n")
	g.printf("\t\tif err != nil {\n")
	g.printf("\t\t\treturn 0, err\n")
	g.printf("\t\t}\n")
	g.printf("\t\tif err := decoder.Decode(v.(map[string]interface{})); err != nil {\n")
	g.print("\t\t\treturn nil, fmt.Errorf(\"Error decoding %+v into %+v: %s\", v, out, err)\n")
	g.printf("\t\t}\n")
	g.printf("\t\treturn out, nil\n")
	g.printf("\t}\n")
	g.printf("\treturn v, nil\n")
	g.printf("}\n")
}

func genDecoderTypeMap(g *generator) {
	g.printf("var objectTypesByTypename = map[string]reflect.Type{\n")
	for _, def := range g.doc.Definitions {
		switch def := def.(type) {
		case *ast.ObjectDefinition:
			g.printf("\"%s\": reflect.TypeOf(%s{}),\n", def.Name.Value, def.Name.Value)
		}
	}
	g.printf("}\n")
}

func genOutputTypeVar(g *generator, f *ast.FieldDefinition) {
	outType := g.goType(f.Type, "")
	if outType[0] == '*' {
		outType = outType[1:]
	}
	g.printf("\tvar out %s\n", outType)
}

func outputTypeReturn(g *generator, f *ast.FieldDefinition) string {
	outType := g.goType(f.Type, "")
	if outType[0] == '*' {
		return "&out"
	}
	return "out"
}

func genClientMethodForField(g *generator, d *ast.ObjectDefinition, f *ast.FieldDefinition) {
	g.printf("func (c *client) %s {\n", renderSignatureForField(g, d, f))
	g.printf("\tc.log.Debugf(ctx, \"%s%s\")\n", exportedName(d.Name.Value), exportedName(f.Name.Value))
	genOutputTypeVar(g, f)
	g.printf("\tif _, err := c.do(ctx, \"%s\", query, &out); err != nil {\n", unexportedName(f.Name.Value))
	g.printf("\t\treturn %s, err\n", renderDefaultReturnValue(f.Type))
	g.printf("\t}\n")
	g.printf("\treturn %s, nil", outputTypeReturn(g, f))
	g.printf("}\n")
}

func renderDefaultClientReturnValue(t ast.Type) string {
	switch t := t.(type) {
	case *ast.NonNull:
		return renderDefaultClientReturnValue(t.Type)
	case *ast.Named:
		switch t.Name.Value {
		case "ID", "String":
			return `""`
		case "Boolean":
			return "false"
		case "Float", "Int":
			return "0"
		}
	}
	return "nil"
}

func genQueryWrapperTypes(g *generator) {
	g.print(`
		type gqlRequestBody struct {
			Query string ` + "`json:\"query\"`" + `
		}

		type gqlResponse struct {
			Data   map[string]interface{}   ` + "`json:\"data\"`" + `
			Errors []map[string]interface{} ` + "`json:\"errors\"`" + `
		}
	`)
}

func genRewriteQuery(g *generator) {
	g.print(`
		func rewriteQuery(query string) (string, error) {
			qast, err := parser.Parse(parser.ParseParams{
				Source: source.New("GraphQL Query", query),
			})
			if err != nil {
				return "", err
			}
			graphql.RequestTypeNames(qast)
			return printer.Print(qast), nil
		}
	`)
}

func genClientDo(g *generator) {
	g.print(`
		func (c *client) do(ctx context.Context, dataField, query string, out interface{}) (int, error) {
			query, err := rewriteQuery(query)
			if err != nil {
				return 0, err
			}
			rb := &gqlRequestBody{Query: query}
			bBody, err := json.Marshal(rb)
			if err != nil {
				return 0, err
			}
			c.log.Debugf(ctx, "Request: %s - %s", c.endpoint, string(bBody))
			req, err := http.NewRequest(http.MethodPost, c.endpoint, bytes.NewReader(bBody))
			if err != nil {
				return 0, err
			}
			if c.authToken != "" {
				req.AddCookie(&http.Cookie{
					Name:     "at",
					Value:    c.authToken,
				})
			}
			for h, hvs := range c.headers {
				for _, hv := range hvs {
					req.Header.Add(h, hv)
				}
			}
			resp, err := http.DefaultClient.Do(req)
			if err != nil {
				return 0, err
			}
			defer resp.Body.Close()
			ball, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return resp.StatusCode, fmt.Errorf("Error reading body - in response from %v: %s", req.URL.String(), err)
			}
			c.log.Debugf(ctx, "Response: %s - %s", resp.Status, ball)
			gqlResp := &gqlResponse{}
			if resp.StatusCode == http.StatusOK {
				if err := json.NewDecoder(bytes.NewReader(ball)).Decode(gqlResp); err != nil {
					return resp.StatusCode, fmt.Errorf("Error parsing body into output: %s", err)
				}
				if len(gqlResp.Errors) != 0 {
					var allErrors string
					for i, err := range gqlResp.Errors {
						allErrors += fmt.Sprintf("%d. - %+v\n", i, err)
					}
					return resp.StatusCode, errors.New(allErrors)
				}
				decoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{
					DecodeHook: decoderHook,
					Result:     out,
				})
				if err != nil {
					return 0, err
				}
				mData, ok := gqlResp.Data[dataField].(map[string]interface{})
				if ok {
					if err := decoder.Decode(mData); err != nil {
						return resp.StatusCode, fmt.Errorf("Error parsing body into output: %s", err)
					}
				} else {
					sData, ok := gqlResp.Data[dataField].([]interface{})
					if ok {
						if err := decoder.Decode(sData); err != nil {
							return resp.StatusCode, fmt.Errorf("Error parsing body into output: %s", err)
						}
					} else {
						return resp.StatusCode, fmt.Errorf("unhandled response data type %T %+v", gqlResp.Data[dataField], gqlResp.Data[dataField])
					}
				}

				return resp.StatusCode, nil
			}
			return resp.StatusCode, fmt.Errorf("Non 200 Response (%d) from %s: %s", resp.StatusCode, req.URL, string(ball))
		}
	`)
}
<|endoftext|>


<commit_before>/*
Copyright The Helm Authors.

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

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

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

package main

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"

	"helm.sh/helm/v3/pkg/repo/repotest"
)

func TestInstall(t *testing.T) {
	srv, err := repotest.NewTempServerWithCleanup(t, "testdata/testcharts/*.tgz*")
	if err != nil {
		t.Fatal(err)
	}
	defer srv.Stop()

	srv.WithMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		username, password, ok := r.BasicAuth()
		if !ok || username != "username" || password != "password" {
			t.Errorf("Expected request to use basic auth and for username == 'username' and password == 'password', got '%v', '%s', '%s'", ok, username, password)
		}
	}))

	srv2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.FileServer(http.Dir(srv.Root())).ServeHTTP(w, r)
	}))
	defer srv2.Close()

	if err := srv.LinkIndices(); err != nil {
		t.Fatal(err)
	}

	tests := []cmdTestCase{
		// Install, base case
		{
			name:   "basic install",
			cmd:    "install aeneas testdata/testcharts/empty --namespace default",
			golden: "output/install.txt",
		},

		// Install, values from cli
		{
			name:   "install with values",
			cmd:    "install virgil testdata/testcharts/alpine --set test.Name=bar",
			golden: "output/install-with-values.txt",
		},
		// Install, values from cli via multiple --set
		{
			name:   "install with multiple values",
			cmd:    "install virgil testdata/testcharts/alpine --set test.Color=yellow --set test.Name=banana",
			golden: "output/install-with-multiple-values.txt",
		},
		// Install, values from yaml
		{
			name:   "install with values file",
			cmd:    "install virgil testdata/testcharts/alpine -f testdata/testcharts/alpine/extra_values.yaml",
			golden: "output/install-with-values-file.txt",
		},
		// Install, no hooks
		{
			name:   "install without hooks",
			cmd:    "install aeneas testdata/testcharts/alpine --no-hooks --set test.Name=hello",
			golden: "output/install-no-hooks.txt",
		},
		// Install, values from multiple yaml
		{
			name:   "install with values",
			cmd:    "install virgil testdata/testcharts/alpine -f testdata/testcharts/alpine/extra_values.yaml -f testdata/testcharts/alpine/more_values.yaml",
			golden: "output/install-with-multiple-values-files.txt",
		},
		// Install, no charts
		{
			name:      "install with no chart specified",
			cmd:       "install",
			golden:    "output/install-no-args.txt",
			wantError: true,
		},
		// Install, re-use name
		{
			name:   "install and replace release",
			cmd:    "install aeneas testdata/testcharts/empty --replace",
			golden: "output/install-and-replace.txt",
		},
		// Install, with timeout
		{
			name:   "install with a timeout",
			cmd:    "install foobar testdata/testcharts/empty --timeout 120s",
			golden: "output/install-with-timeout.txt",
		},
		// Install, with wait
		{
			name:   "install with a wait",
			cmd:    "install apollo testdata/testcharts/empty --wait",
			golden: "output/install-with-wait.txt",
		},
		// Install, with wait-for-jobs
		{
			name:   "install with wait-for-jobs",
			cmd:    "install apollo testdata/testcharts/empty --wait --wait-for-jobs",
			golden: "output/install-with-wait-for-jobs.txt",
		},
		// Install, using the name-template
		{
			name:   "install with name-template",
			cmd:    "install testdata/testcharts/empty --name-template '{{upper \"foobar\"}}'",
			golden: "output/install-name-template.txt",
		},
		// Install, perform chart verification along the way.
		{
			name:      "install with verification, missing provenance",
			cmd:       "install bogus testdata/testcharts/compressedchart-0.1.0.tgz --verify --keyring testdata/helm-test-key.pub",
			wantError: true,
		},
		{
			name:      "install with verification, directory instead of file",
			cmd:       "install bogus testdata/testcharts/signtest --verify --keyring testdata/helm-test-key.pub",
			wantError: true,
		},
		{
			name: "install with verification, valid",
			cmd:  "install signtest testdata/testcharts/signtest-0.1.0.tgz --verify --keyring testdata/helm-test-key.pub",
		},
		// Install, chart with missing dependencies in /charts
		{
			name:      "install chart with missing dependencies",
			cmd:       "install nodeps testdata/testcharts/chart-missing-deps",
			wantError: true,
		},
		// Install chart with update-dependency
		{
			name:   "install chart with missing dependencies",
			cmd:    "install --dependency-update updeps testdata/testcharts/chart-with-subchart-update",
			golden: "output/chart-with-subchart-update.txt",
		},
		// Install, chart with bad dependencies in Chart.yaml in /charts
		{
			name:      "install chart with bad dependencies in Chart.yaml",
			cmd:       "install badreq testdata/testcharts/chart-bad-requirements",
			wantError: true,
		},
		// Install, chart with library chart dependency
		{
			name: "install chart with library chart dependency",
			cmd:  "install withlibchartp testdata/testcharts/chart-with-lib-dep",
		},
		// Install, library chart
		{
			name:      "install library chart",
			cmd:       "install libchart testdata/testcharts/lib-chart",
			wantError: true,
			golden:    "output/template-lib-chart.txt",
		},
		// Install, chart with bad type
		{
			name:      "install chart with bad type",
			cmd:       "install badtype testdata/testcharts/chart-bad-type",
			wantError: true,
			golden:    "output/install-chart-bad-type.txt",
		},
		// Install, values from yaml, schematized
		{
			name:   "install with schema file",
			cmd:    "install schema testdata/testcharts/chart-with-schema",
			golden: "output/schema.txt",
		},
		// Install, values from yaml, schematized with errors
		{
			name:      "install with schema file, with errors",
			cmd:       "install schema testdata/testcharts/chart-with-schema-negative",
			wantError: true,
			golden:    "output/schema-negative.txt",
		},
		// Install, values from yaml, extra values from yaml, schematized with errors
		{
			name:      "install with schema file, extra values from yaml, with errors",
			cmd:       "install schema testdata/testcharts/chart-with-schema -f testdata/testcharts/chart-with-schema/extra-values.yaml",
			wantError: true,
			golden:    "output/schema-negative.txt",
		},
		// Install, values from yaml, extra values from cli, schematized with errors
		{
			name:      "install with schema file, extra values from cli, with errors",
			cmd:       "install schema testdata/testcharts/chart-with-schema --set age=-5",
			wantError: true,
			golden:    "output/schema-negative-cli.txt",
		},
		// Install with subchart, values from yaml, schematized with errors
		{
			name:      "install with schema file and schematized subchart, with errors",
			cmd:       "install schema testdata/testcharts/chart-with-schema-and-subchart",
			wantError: true,
			golden:    "output/subchart-schema-negative.txt",
		},
		// Install with subchart, values from yaml, extra values from cli, schematized with errors
		{
			name:   "install with schema file and schematized subchart, extra values from cli",
			cmd:    "install schema testdata/testcharts/chart-with-schema-and-subchart --set lastname=doe --set subchart-with-schema.age=25",
			golden: "output/subchart-schema-cli.txt",
		},
		// Install with subchart, values from yaml, extra values from cli, schematized with errors
		{
			name:      "install with schema file and schematized subchart, extra values from cli, with errors",
			cmd:       "install schema testdata/testcharts/chart-with-schema-and-subchart --set lastname=doe --set subchart-with-schema.age=-25",
			wantError: true,
			golden:    "output/subchart-schema-cli-negative.txt",
		},
		// Install deprecated chart
		{
			name:   "install with warning about deprecated chart",
			cmd:    "install aeneas testdata/testcharts/deprecated --namespace default",
			golden: "output/deprecated-chart.txt",
		},
		// Install chart with only crds
		{
			name: "install chart with only crds",
			cmd:  "install crd-test testdata/testcharts/chart-with-only-crds --namespace default",
		},
		// Verify the user/pass works
		{
			name:   "basic install with credentials",
			cmd:    "install aeneas reqtest --namespace default --repo " + srv.URL() + " --username username --password password",
			golden: "output/install.txt",
		},
		{
			name:   "basic install with credentials",
			cmd:    "install aeneas reqtest --namespace default --repo " + srv2.URL + " --username username --password password --pass-credentials",
			golden: "output/install.txt",
		},
	}

	runTestActionCmd(t, tests)
}

func TestInstallOutputCompletion(t *testing.T) {
	outputFlagCompletionTest(t, "install")
}

func TestInstallVersionCompletion(t *testing.T) {
	repoFile := "testdata/helmhome/helm/repositories.yaml"
	repoCache := "testdata/helmhome/helm/repository"

	repoSetup := fmt.Sprintf("--repository-config %s --repository-cache %s", repoFile, repoCache)

	tests := []cmdTestCase{{
		name:   "completion for install version flag with release name",
		cmd:    fmt.Sprintf("%s __complete install releasename testing/alpine --version ''", repoSetup),
		golden: "output/version-comp.txt",
	}, {
		name:   "completion for install version flag with generate-name",
		cmd:    fmt.Sprintf("%s __complete install --generate-name testing/alpine --version ''", repoSetup),
		golden: "output/version-comp.txt",
	}, {
		name:   "completion for install version flag too few args",
		cmd:    fmt.Sprintf("%s __complete install testing/alpine --version ''", repoSetup),
		golden: "output/version-invalid-comp.txt",
	}, {
		name:   "completion for install version flag too many args",
		cmd:    fmt.Sprintf("%s __complete install releasename testing/alpine badarg --version ''", repoSetup),
		golden: "output/version-invalid-comp.txt",
	}, {
		name:   "completion for install version flag invalid chart",
		cmd:    fmt.Sprintf("%s __complete install releasename invalid/invalid --version ''", repoSetup),
		golden: "output/version-invalid-comp.txt",
	}}
	runTestCmd(t, tests)
}

func TestInstallFileCompletion(t *testing.T) {
	checkFileCompletion(t, "install", false)
	checkFileCompletion(t, "install --generate-name", true)
	checkFileCompletion(t, "install myname", true)
	checkFileCompletion(t, "install myname mychart", false)
}
<commit_msg>Adding test for user/pass without repo on Helm install<commit_after>/*
Copyright The Helm Authors.

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

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

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

package main

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"testing"

	"helm.sh/helm/v3/pkg/repo/repotest"
)

func TestInstall(t *testing.T) {
	srv, err := repotest.NewTempServerWithCleanup(t, "testdata/testcharts/*.tgz*")
	if err != nil {
		t.Fatal(err)
	}
	defer srv.Stop()

	srv.WithMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		username, password, ok := r.BasicAuth()
		if !ok || username != "username" || password != "password" {
			t.Errorf("Expected request to use basic auth and for username == 'username' and password == 'password', got '%v', '%s', '%s'", ok, username, password)
		}
	}))

	srv2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.FileServer(http.Dir(srv.Root())).ServeHTTP(w, r)
	}))
	defer srv2.Close()

	if err := srv.LinkIndices(); err != nil {
		t.Fatal(err)
	}

	repoFile := filepath.Join(srv.Root(), "repositories.yaml")

	tests := []cmdTestCase{
		// Install, base case
		{
			name:   "basic install",
			cmd:    "install aeneas testdata/testcharts/empty --namespace default",
			golden: "output/install.txt",
		},

		// Install, values from cli
		{
			name:   "install with values",
			cmd:    "install virgil testdata/testcharts/alpine --set test.Name=bar",
			golden: "output/install-with-values.txt",
		},
		// Install, values from cli via multiple --set
		{
			name:   "install with multiple values",
			cmd:    "install virgil testdata/testcharts/alpine --set test.Color=yellow --set test.Name=banana",
			golden: "output/install-with-multiple-values.txt",
		},
		// Install, values from yaml
		{
			name:   "install with values file",
			cmd:    "install virgil testdata/testcharts/alpine -f testdata/testcharts/alpine/extra_values.yaml",
			golden: "output/install-with-values-file.txt",
		},
		// Install, no hooks
		{
			name:   "install without hooks",
			cmd:    "install aeneas testdata/testcharts/alpine --no-hooks --set test.Name=hello",
			golden: "output/install-no-hooks.txt",
		},
		// Install, values from multiple yaml
		{
			name:   "install with values",
			cmd:    "install virgil testdata/testcharts/alpine -f testdata/testcharts/alpine/extra_values.yaml -f testdata/testcharts/alpine/more_values.yaml",
			golden: "output/install-with-multiple-values-files.txt",
		},
		// Install, no charts
		{
			name:      "install with no chart specified",
			cmd:       "install",
			golden:    "output/install-no-args.txt",
			wantError: true,
		},
		// Install, re-use name
		{
			name:   "install and replace release",
			cmd:    "install aeneas testdata/testcharts/empty --replace",
			golden: "output/install-and-replace.txt",
		},
		// Install, with timeout
		{
			name:   "install with a timeout",
			cmd:    "install foobar testdata/testcharts/empty --timeout 120s",
			golden: "output/install-with-timeout.txt",
		},
		// Install, with wait
		{
			name:   "install with a wait",
			cmd:    "install apollo testdata/testcharts/empty --wait",
			golden: "output/install-with-wait.txt",
		},
		// Install, with wait-for-jobs
		{
			name:   "install with wait-for-jobs",
			cmd:    "install apollo testdata/testcharts/empty --wait --wait-for-jobs",
			golden: "output/install-with-wait-for-jobs.txt",
		},
		// Install, using the name-template
		{
			name:   "install with name-template",
			cmd:    "install testdata/testcharts/empty --name-template '{{upper \"foobar\"}}'",
			golden: "output/install-name-template.txt",
		},
		// Install, perform chart verification along the way.
		{
			name:      "install with verification, missing provenance",
			cmd:       "install bogus testdata/testcharts/compressedchart-0.1.0.tgz --verify --keyring testdata/helm-test-key.pub",
			wantError: true,
		},
		{
			name:      "install with verification, directory instead of file",
			cmd:       "install bogus testdata/testcharts/signtest --verify --keyring testdata/helm-test-key.pub",
			wantError: true,
		},
		{
			name: "install with verification, valid",
			cmd:  "install signtest testdata/testcharts/signtest-0.1.0.tgz --verify --keyring testdata/helm-test-key.pub",
		},
		// Install, chart with missing dependencies in /charts
		{
			name:      "install chart with missing dependencies",
			cmd:       "install nodeps testdata/testcharts/chart-missing-deps",
			wantError: true,
		},
		// Install chart with update-dependency
		{
			name:   "install chart with missing dependencies",
			cmd:    "install --dependency-update updeps testdata/testcharts/chart-with-subchart-update",
			golden: "output/chart-with-subchart-update.txt",
		},
		// Install, chart with bad dependencies in Chart.yaml in /charts
		{
			name:      "install chart with bad dependencies in Chart.yaml",
			cmd:       "install badreq testdata/testcharts/chart-bad-requirements",
			wantError: true,
		},
		// Install, chart with library chart dependency
		{
			name: "install chart with library chart dependency",
			cmd:  "install withlibchartp testdata/testcharts/chart-with-lib-dep",
		},
		// Install, library chart
		{
			name:      "install library chart",
			cmd:       "install libchart testdata/testcharts/lib-chart",
			wantError: true,
			golden:    "output/template-lib-chart.txt",
		},
		// Install, chart with bad type
		{
			name:      "install chart with bad type",
			cmd:       "install badtype testdata/testcharts/chart-bad-type",
			wantError: true,
			golden:    "output/install-chart-bad-type.txt",
		},
		// Install, values from yaml, schematized
		{
			name:   "install with schema file",
			cmd:    "install schema testdata/testcharts/chart-with-schema",
			golden: "output/schema.txt",
		},
		// Install, values from yaml, schematized with errors
		{
			name:      "install with schema file, with errors",
			cmd:       "install schema testdata/testcharts/chart-with-schema-negative",
			wantError: true,
			golden:    "output/schema-negative.txt",
		},
		// Install, values from yaml, extra values from yaml, schematized with errors
		{
			name:      "install with schema file, extra values from yaml, with errors",
			cmd:       "install schema testdata/testcharts/chart-with-schema -f testdata/testcharts/chart-with-schema/extra-values.yaml",
			wantError: true,
			golden:    "output/schema-negative.txt",
		},
		// Install, values from yaml, extra values from cli, schematized with errors
		{
			name:      "install with schema file, extra values from cli, with errors",
			cmd:       "install schema testdata/testcharts/chart-with-schema --set age=-5",
			wantError: true,
			golden:    "output/schema-negative-cli.txt",
		},
		// Install with subchart, values from yaml, schematized with errors
		{
			name:      "install with schema file and schematized subchart, with errors",
			cmd:       "install schema testdata/testcharts/chart-with-schema-and-subchart",
			wantError: true,
			golden:    "output/subchart-schema-negative.txt",
		},
		// Install with subchart, values from yaml, extra values from cli, schematized with errors
		{
			name:   "install with schema file and schematized subchart, extra values from cli",
			cmd:    "install schema testdata/testcharts/chart-with-schema-and-subchart --set lastname=doe --set subchart-with-schema.age=25",
			golden: "output/subchart-schema-cli.txt",
		},
		// Install with subchart, values from yaml, extra values from cli, schematized with errors
		{
			name:      "install with schema file and schematized subchart, extra values from cli, with errors",
			cmd:       "install schema testdata/testcharts/chart-with-schema-and-subchart --set lastname=doe --set subchart-with-schema.age=-25",
			wantError: true,
			golden:    "output/subchart-schema-cli-negative.txt",
		},
		// Install deprecated chart
		{
			name:   "install with warning about deprecated chart",
			cmd:    "install aeneas testdata/testcharts/deprecated --namespace default",
			golden: "output/deprecated-chart.txt",
		},
		// Install chart with only crds
		{
			name: "install chart with only crds",
			cmd:  "install crd-test testdata/testcharts/chart-with-only-crds --namespace default",
		},
		// Verify the user/pass works
		{
			name:   "basic install with credentials",
			cmd:    "install aeneas reqtest --namespace default --repo " + srv.URL() + " --username username --password password",
			golden: "output/install.txt",
		},
		{
			name:   "basic install with credentials",
			cmd:    "install aeneas reqtest --namespace default --repo " + srv2.URL + " --username username --password password --pass-credentials",
			golden: "output/install.txt",
		},
		{
			name:   "basic install with credentials and no repo",
			cmd:    fmt.Sprintf("install aeneas test/reqtest --username username --password password --repository-config %s --repository-cache %s", repoFile, srv.Root()),
			golden: "output/install.txt",
		},
	}

	runTestActionCmd(t, tests)
}

func TestInstallOutputCompletion(t *testing.T) {
	outputFlagCompletionTest(t, "install")
}

func TestInstallVersionCompletion(t *testing.T) {
	repoFile := "testdata/helmhome/helm/repositories.yaml"
	repoCache := "testdata/helmhome/helm/repository"

	repoSetup := fmt.Sprintf("--repository-config %s --repository-cache %s", repoFile, repoCache)

	tests := []cmdTestCase{{
		name:   "completion for install version flag with release name",
		cmd:    fmt.Sprintf("%s __complete install releasename testing/alpine --version ''", repoSetup),
		golden: "output/version-comp.txt",
	}, {
		name:   "completion for install version flag with generate-name",
		cmd:    fmt.Sprintf("%s __complete install --generate-name testing/alpine --version ''", repoSetup),
		golden: "output/version-comp.txt",
	}, {
		name:   "completion for install version flag too few args",
		cmd:    fmt.Sprintf("%s __complete install testing/alpine --version ''", repoSetup),
		golden: "output/version-invalid-comp.txt",
	}, {
		name:   "completion for install version flag too many args",
		cmd:    fmt.Sprintf("%s __complete install releasename testing/alpine badarg --version ''", repoSetup),
		golden: "output/version-invalid-comp.txt",
	}, {
		name:   "completion for install version flag invalid chart",
		cmd:    fmt.Sprintf("%s __complete install releasename invalid/invalid --version ''", repoSetup),
		golden: "output/version-invalid-comp.txt",
	}}
	runTestCmd(t, tests)
}

func TestInstallFileCompletion(t *testing.T) {
	checkFileCompletion(t, "install", false)
	checkFileCompletion(t, "install --generate-name", true)
	checkFileCompletion(t, "install myname", true)
	checkFileCompletion(t, "install myname mychart", false)
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/fd0/khepri"
	"github.com/fd0/khepri/backend"
	"golang.org/x/crypto/ssh/terminal"
)

func format_bytes(c uint64) string {
	b := float64(c)

	switch {
	case c > 1<<40:
		return fmt.Sprintf("%.3fTiB", b/(1<<40))
	case c > 1<<30:
		return fmt.Sprintf("%.3fGiB", b/(1<<30))
	case c > 1<<20:
		return fmt.Sprintf("%.3fMiB", b/(1<<20))
	case c > 1<<10:
		return fmt.Sprintf("%.3fKiB", b/(1<<10))
	default:
		return fmt.Sprintf("%dB", c)
	}
}

func format_duration(sec uint64) string {
	hours := sec / 3600
	sec -= hours * 3600
	min := sec / 60
	sec -= min * 60
	if hours > 0 {
		return fmt.Sprintf("%d:%02d:%02d", hours, min, sec)
	}

	return fmt.Sprintf("%d:%02d", min, sec)
}

func print_tree2(indent int, t *khepri.Tree) {
	for _, node := range *t {
		if node.Tree != nil {
			fmt.Printf("%s%s/\n", strings.Repeat("  ", indent), node.Name)
			print_tree2(indent+1, node.Tree)
		} else {
			fmt.Printf("%s%s\n", strings.Repeat("  ", indent), node.Name)
		}
	}
}

func commandBackup(be backend.Server, key *khepri.Key, args []string) error {
	if len(args) != 1 {
		return errors.New("usage: backup [dir|file]")
	}

	target := args[0]

	arch, err := khepri.NewArchiver(be, key)
	if err != nil {
		fmt.Fprintf(os.Stderr, "err: %v\n", err)
	}
	arch.Error = func(dir string, fi os.FileInfo, err error) error {
		// TODO: make ignoring errors configurable
		fmt.Fprintf(os.Stderr, "\nerror for %s: %v\n%v\n", dir, err, fi)
		return nil
	}

	fmt.Printf("scanning %s\n", target)

	if terminal.IsTerminal(int(os.Stdout.Fd())) {
		ch := make(chan khepri.Stats, 20)
		arch.ScannerStats = ch

		go func(ch <-chan khepri.Stats) {
			for stats := range ch {
				fmt.Printf("\r%6d directories, %6d files, %14s", stats.Directories, stats.Files, format_bytes(stats.Bytes))
			}
		}(ch)
	}

	fmt.Printf("done\n")

	// TODO: add filter
	// arch.Filter = func(dir string, fi os.FileInfo) bool {
	// 	return true
	// }

	t, err := arch.LoadTree(target)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		return err
	}

	fmt.Printf("\r%6d directories, %6d files, %14s\n", arch.Stats.Directories, arch.Stats.Files, format_bytes(arch.Stats.Bytes))

	stats := khepri.Stats{}
	start := time.Now()
	if terminal.IsTerminal(int(os.Stdout.Fd())) {
		ch := make(chan khepri.Stats, 20)
		arch.SaveStats = ch

		ticker := time.NewTicker(time.Second)
		var eta, bps uint64

		go func(ch <-chan khepri.Stats) {

			status := func(d time.Duration) {
				fmt.Printf("\x1b[2K\r[%s] %3.2f%%  %s/s  %s / %s  ETA %s",
					format_duration(uint64(d/time.Second)),
					float64(stats.Bytes)/float64(arch.Stats.Bytes)*100,
					format_bytes(bps),
					format_bytes(stats.Bytes), format_bytes(arch.Stats.Bytes),
					format_duration(eta))
			}

			defer ticker.Stop()
			for {
				select {
				case s, ok := <-ch:
					if !ok {
						return
					}
					stats.Files += s.Files
					stats.Directories += s.Directories
					stats.Other += s.Other
					stats.Bytes += s.Bytes

					status(time.Since(start))
				case <-ticker.C:
					d := time.Since(start)
					bps = stats.Bytes * uint64(time.Second) / uint64(d)

					if bps > 0 {
						eta = (arch.Stats.Bytes - stats.Bytes) / bps
					}

					status(d)
				}
			}
		}(ch)
	}

	sn, id, err := arch.Snapshot(target, t)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
	}

	if terminal.IsTerminal(int(os.Stdout.Fd())) {
		// close channels so that the goroutines terminate
		close(arch.SaveStats)
		close(arch.ScannerStats)
	}

	fmt.Printf("\nsnapshot %s saved: %v\n", id, sn)
	duration := time.Now().Sub(start)
	fmt.Printf("duration: %s, %.2fMiB/s\n", duration, float64(arch.Stats.Bytes)/float64(duration/time.Second)/(1<<20))

	return nil
}
<commit_msg>Add space between number and unit<commit_after>package main

import (
	"errors"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/fd0/khepri"
	"github.com/fd0/khepri/backend"
	"golang.org/x/crypto/ssh/terminal"
)

func format_bytes(c uint64) string {
	b := float64(c)

	switch {
	case c > 1<<40:
		return fmt.Sprintf("%.3f TiB", b/(1<<40))
	case c > 1<<30:
		return fmt.Sprintf("%.3f GiB", b/(1<<30))
	case c > 1<<20:
		return fmt.Sprintf("%.3f MiB", b/(1<<20))
	case c > 1<<10:
		return fmt.Sprintf("%.3f KiB", b/(1<<10))
	default:
		return fmt.Sprintf("%dB", c)
	}
}

func format_duration(sec uint64) string {
	hours := sec / 3600
	sec -= hours * 3600
	min := sec / 60
	sec -= min * 60
	if hours > 0 {
		return fmt.Sprintf("%d:%02d:%02d", hours, min, sec)
	}

	return fmt.Sprintf("%d:%02d", min, sec)
}

func print_tree2(indent int, t *khepri.Tree) {
	for _, node := range *t {
		if node.Tree != nil {
			fmt.Printf("%s%s/\n", strings.Repeat("  ", indent), node.Name)
			print_tree2(indent+1, node.Tree)
		} else {
			fmt.Printf("%s%s\n", strings.Repeat("  ", indent), node.Name)
		}
	}
}

func commandBackup(be backend.Server, key *khepri.Key, args []string) error {
	if len(args) != 1 {
		return errors.New("usage: backup [dir|file]")
	}

	target := args[0]

	arch, err := khepri.NewArchiver(be, key)
	if err != nil {
		fmt.Fprintf(os.Stderr, "err: %v\n", err)
	}
	arch.Error = func(dir string, fi os.FileInfo, err error) error {
		// TODO: make ignoring errors configurable
		fmt.Fprintf(os.Stderr, "\nerror for %s: %v\n%v\n", dir, err, fi)
		return nil
	}

	fmt.Printf("scanning %s\n", target)

	if terminal.IsTerminal(int(os.Stdout.Fd())) {
		ch := make(chan khepri.Stats, 20)
		arch.ScannerStats = ch

		go func(ch <-chan khepri.Stats) {
			for stats := range ch {
				fmt.Printf("\r%6d directories, %6d files, %14s", stats.Directories, stats.Files, format_bytes(stats.Bytes))
			}
		}(ch)
	}

	fmt.Printf("done\n")

	// TODO: add filter
	// arch.Filter = func(dir string, fi os.FileInfo) bool {
	// 	return true
	// }

	t, err := arch.LoadTree(target)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		return err
	}

	fmt.Printf("\r%6d directories, %6d files, %14s\n", arch.Stats.Directories, arch.Stats.Files, format_bytes(arch.Stats.Bytes))

	stats := khepri.Stats{}
	start := time.Now()
	if terminal.IsTerminal(int(os.Stdout.Fd())) {
		ch := make(chan khepri.Stats, 20)
		arch.SaveStats = ch

		ticker := time.NewTicker(time.Second)
		var eta, bps uint64

		go func(ch <-chan khepri.Stats) {

			status := func(d time.Duration) {
				fmt.Printf("\x1b[2K\r[%s] %3.2f%%  %s/s  %s / %s  ETA %s",
					format_duration(uint64(d/time.Second)),
					float64(stats.Bytes)/float64(arch.Stats.Bytes)*100,
					format_bytes(bps),
					format_bytes(stats.Bytes), format_bytes(arch.Stats.Bytes),
					format_duration(eta))
			}

			defer ticker.Stop()
			for {
				select {
				case s, ok := <-ch:
					if !ok {
						return
					}
					stats.Files += s.Files
					stats.Directories += s.Directories
					stats.Other += s.Other
					stats.Bytes += s.Bytes

					status(time.Since(start))
				case <-ticker.C:
					d := time.Since(start)
					bps = stats.Bytes * uint64(time.Second) / uint64(d)

					if bps > 0 {
						eta = (arch.Stats.Bytes - stats.Bytes) / bps
					}

					status(d)
				}
			}
		}(ch)
	}

	sn, id, err := arch.Snapshot(target, t)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
	}

	if terminal.IsTerminal(int(os.Stdout.Fd())) {
		// close channels so that the goroutines terminate
		close(arch.SaveStats)
		close(arch.ScannerStats)
	}

	fmt.Printf("\nsnapshot %s saved: %v\n", id, sn)
	duration := time.Now().Sub(start)
	fmt.Printf("duration: %s, %.2fMiB/s\n", duration, float64(arch.Stats.Bytes)/float64(duration/time.Second)/(1<<20))

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package util

import (
	"fmt"
	"net/url"
	"strings"

	"k8s.io/apimachinery/pkg/util/validation/field"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/klog"
	gceacls "k8s.io/kops/pkg/acls/gce"
	s3acls "k8s.io/kops/pkg/acls/s3"
	kopsclient "k8s.io/kops/pkg/client/clientset_generated/clientset"
	"k8s.io/kops/pkg/client/simple"
	"k8s.io/kops/pkg/client/simple/api"
	"k8s.io/kops/pkg/client/simple/vfsclientset"
	"k8s.io/kops/util/pkg/vfs"
)

type FactoryOptions struct {
	RegistryPath string
}

type Factory struct {
	options   *FactoryOptions
	clientset simple.Clientset
}

func NewFactory(options *FactoryOptions) *Factory {
	gceacls.Register()
	s3acls.Register()

	return &Factory{
		options: options,
	}
}

const (
	STATE_ERROR = `Please set the --state flag or export KOPS_STATE_STORE.
For example, a valid value follows the format s3://<bucket>.
You can find the supported stores in https://kops.sigs.k8s.io/state.`

	INVALID_STATE_ERROR = `Unable to read state store.
Please use a valid state store when setting --state or KOPS_STATE_STORE env var.
For example, a valid value follows the format s3://<bucket>.
Trailing slash will be trimmed.`
)

func (f *Factory) Clientset() (simple.Clientset, error) {
	if f.clientset == nil {
		registryPath := f.options.RegistryPath
		klog.V(2).Infof("state store %s", registryPath)
		if registryPath == "" {
			return nil, field.Required(field.NewPath("State Store"), STATE_ERROR)
		}

		// We recognize a `k8s` scheme; this might change in future so we won't document it yet
		// In practice nobody is going to hit this accidentally, so I don't think we need a feature flag.
		if strings.HasPrefix(registryPath, "k8s://") {
			loadingRules := clientcmd.NewDefaultClientConfigLoadingRules()

			configOverrides := &clientcmd.ConfigOverrides{}

			if registryPath == "k8s://" {
			} else {
				u, err := url.Parse(registryPath)
				if err != nil {
					return nil, fmt.Errorf("invalid kops server url: %q", registryPath)
				}
				configOverrides.CurrentContext = u.Host
			}

			kubeConfig := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(loadingRules, configOverrides)
			config, err := kubeConfig.ClientConfig()
			if err != nil {
				return nil, fmt.Errorf("error loading kubeconfig for %q", registryPath)
			}

			kopsClient, err := kopsclient.NewForConfig(config)
			if err != nil {
				return nil, fmt.Errorf("error building kops API client: %v", err)
			}

			f.clientset = &api.RESTClientset{
				BaseURL: &url.URL{
					Scheme: "k8s",
				},
				KopsClient: kopsClient.Kops(),
			}
		} else {
			basePath, err := vfs.Context.BuildVfsPath(registryPath)
			if err != nil {
				return nil, fmt.Errorf("error building path for %q: %v", registryPath, err)
			}

			if !vfs.IsClusterReadable(basePath) {
				return nil, field.Invalid(field.NewPath("State Store"), registryPath, INVALID_STATE_ERROR)
			}

			// For kops CLI / controller, we do allow vfs list (unlike nodeup!)
			allowVFSList := true

			f.clientset = vfsclientset.NewVFSClientset(basePath, allowVFSList)
		}
	}

	return f.clientset, nil
}
<commit_msg>Add a warning when using file:// state store<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package util

import (
	"fmt"
	"net/url"
	"strings"

	"k8s.io/apimachinery/pkg/util/validation/field"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/klog"
	gceacls "k8s.io/kops/pkg/acls/gce"
	s3acls "k8s.io/kops/pkg/acls/s3"
	kopsclient "k8s.io/kops/pkg/client/clientset_generated/clientset"
	"k8s.io/kops/pkg/client/simple"
	"k8s.io/kops/pkg/client/simple/api"
	"k8s.io/kops/pkg/client/simple/vfsclientset"
	"k8s.io/kops/util/pkg/vfs"
)

type FactoryOptions struct {
	RegistryPath string
}

type Factory struct {
	options   *FactoryOptions
	clientset simple.Clientset
}

func NewFactory(options *FactoryOptions) *Factory {
	gceacls.Register()
	s3acls.Register()

	return &Factory{
		options: options,
	}
}

const (
	STATE_ERROR = `Please set the --state flag or export KOPS_STATE_STORE.
For example, a valid value follows the format s3://<bucket>.
You can find the supported stores in https://kops.sigs.k8s.io/state.`

	INVALID_STATE_ERROR = `Unable to read state store.
Please use a valid state store when setting --state or KOPS_STATE_STORE env var.
For example, a valid value follows the format s3://<bucket>.
Trailing slash will be trimmed.`
)

func (f *Factory) Clientset() (simple.Clientset, error) {
	if f.clientset == nil {
		registryPath := f.options.RegistryPath
		klog.V(2).Infof("state store %s", registryPath)
		if registryPath == "" {
			return nil, field.Required(field.NewPath("State Store"), STATE_ERROR)
		}

		// We recognize a `k8s` scheme; this might change in future so we won't document it yet
		// In practice nobody is going to hit this accidentally, so I don't think we need a feature flag.
		if strings.HasPrefix(registryPath, "k8s://") {
			loadingRules := clientcmd.NewDefaultClientConfigLoadingRules()

			configOverrides := &clientcmd.ConfigOverrides{}

			if registryPath == "k8s://" {
			} else {
				u, err := url.Parse(registryPath)
				if err != nil {
					return nil, fmt.Errorf("invalid kops server url: %q", registryPath)
				}
				configOverrides.CurrentContext = u.Host
			}

			kubeConfig := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(loadingRules, configOverrides)
			config, err := kubeConfig.ClientConfig()
			if err != nil {
				return nil, fmt.Errorf("error loading kubeconfig for %q", registryPath)
			}

			kopsClient, err := kopsclient.NewForConfig(config)
			if err != nil {
				return nil, fmt.Errorf("error building kops API client: %v", err)
			}

			f.clientset = &api.RESTClientset{
				BaseURL: &url.URL{
					Scheme: "k8s",
				},
				KopsClient: kopsClient.Kops(),
			}
		} else {
			basePath, err := vfs.Context.BuildVfsPath(registryPath)
			if err != nil {
				return nil, fmt.Errorf("error building path for %q: %v", registryPath, err)
			}

			if !vfs.IsClusterReadable(basePath) {
				return nil, field.Invalid(field.NewPath("State Store"), registryPath, INVALID_STATE_ERROR)
			}

			// For kops CLI / controller, we do allow vfs list (unlike nodeup!)
			allowVFSList := true

			f.clientset = vfsclientset.NewVFSClientset(basePath, allowVFSList)
		}
		if strings.HasPrefix(registryPath, "file://") {
			klog.Warning("The local filesystem state store is not functional for running clusters")
		}
	}

	return f.clientset, nil
}
<|endoftext|>