add helper functions for running prometheus query

Author: sallyom

pkg/metrics/query.go

@@ -0,0 +1,254 @@
+package metrics
+
+import (
+	"crypto/tls"
+	"crypto/x509"
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"os"
+	"strings"
+	"time"
+
+	corev1 "k8s.io/api/core/v1"
+	kerrs "k8s.io/apimachinery/pkg/api/errors"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/util/wait"
+	"k8s.io/client-go/kubernetes"
+
+	routeclient "github.com/openshift/client-go/route/clientset/versioned"
+	"github.com/prometheus/common/model"
+)
+
+const (
+	serviceCAinjectionDataKey             = "service-ca.crt"
+	serviceCAinjectCABundleAnnotationName = "service.beta.openshift.io/inject-cabundle"
+)
+
+func LocatePrometheus(client *kubernetes.Clientset, rc *routeclient.Clientset) (string, string, error) {
+	_, err := client.CoreV1().Services("openshift-monitoring").Get("prometheus-k8s", metav1.GetOptions{})
+	if err != nil {
+		return "", "", err
+	}
+
+	route, err := rc.RouteV1().Routes("openshift-monitoring").Get("prometheus-k8s", metav1.GetOptions{})
+	if err != nil {
+		return "", "", err
+	}
+	host := route.Spec.Host
+	var bearerToken string
+	secrets, err := client.CoreV1().Secrets("openshift-monitoring").List(metav1.ListOptions{})
+	if err != nil {
+		return "", "", fmt.Errorf("could not list secrets in openshift-monitoring namespace")
+	}
+	for _, secret := range secrets.Items {
+		if secret.Type != corev1.SecretTypeServiceAccountToken {
+			continue
+		}
+		if !strings.HasPrefix(secret.Name, "prometheus-") {
+			continue
+		}
+		bearerToken = string(secret.Data[corev1.ServiceAccountTokenKey])
+		break
+	}
+	if len(bearerToken) == 0 {
+		return "", "", fmt.Errorf("error getting bearer token from prometheus service account")
+	}
+	return bearerToken, host, nil
+}
+
+type prometheusResponse struct {
+	Status string                 `json:"status"`
+	Data   prometheusResponseData `json:"data"`
+}
+
+type prometheusResponseData struct {
+	ResultType string       `json:"resultType"`
+	Result     model.Vector `json:"result"`
+}
+
+// RunPrometheusQuery() checks that query returns non-empty result, errors if query returns empty, and returns metrics
+// example: query := `ALERTS{alertstate="pending",alertname="PodDisruptionBudgetAtLimit",severity="warning"} == 1`
+func RunQuery(caBundle []byte, tlsCert tls.Certificate, host, query, bearerToken string) (model.Vector, error) {
+	var metrics model.Vector
+	err := wait.PollImmediate(time.Second*3, time.Second*300, func() (bool, error) {
+		certPool := x509.NewCertPool()
+		if ok := certPool.AppendCertsFromPEM(caBundle); !ok {
+			return false, fmt.Errorf("error appending caBundle to pool")
+		}
+
+		tlsConfig := &tls.Config{
+			InsecureSkipVerify: true,
+			RootCAs:            certPool,
+			Certificates:       []tls.Certificate{tlsCert},
+		}
+		tlsConfig.BuildNameToCertificate()
+		tr := &http.Transport{TLSClientConfig: tlsConfig}
+		prometheusClient := &http.Client{Transport: tr}
+		req, err := http.NewRequest("GET", "https://"+host+"/api/v1/query", nil)
+		if err != nil {
+			return false, err
+		}
+
+		q := req.URL.Query()
+		q.Add("query", query)
+		req.URL.RawQuery = q.Encode()
+
+		req.Header.Add("Authorization", "Bearer "+bearerToken)
+
+		resp, err := prometheusClient.Do(req)
+		if err != nil {
+			return false, fmt.Errorf("prometheus query failed: %v\n%v", err, resp)
+		}
+		defer resp.Body.Close()
+		if resp.StatusCode != http.StatusOK {
+			return false, fmt.Errorf("Non-200 response from prometheus: %v", resp)
+		}
+
+		body, err := ioutil.ReadAll(resp.Body)
+		result := prometheusResponse{}
+		json.Unmarshal(body, &result)
+		metrics = result.Data.Result
+		if len(metrics) == 0 {
+			return false, fmt.Errorf("prometheus returned unexpected results")
+		}
+		return true, nil
+	})
+	if err != nil {
+		return nil, fmt.Errorf("prometheus query error: %v", err)
+	}
+	return metrics, nil
+}
+
+func RetrieveTLSCert(client *kubernetes.Clientset, secret, svc, ns string) (tls.Certificate, error) {
+	err := createServingCertAnnotatedService(client, secret, svc, ns)
+	if err != nil {
+		return tls.Certificate{}, err
+	}
+	err = pollForServiceServingSecret(client, secret, ns)
+	if err != nil {
+		return tls.Certificate{}, err
+	}
+	cert, err := getServiceServingCertSecretData(client, secret, ns)
+	if err != nil {
+		return tls.Certificate{}, err
+	}
+	return cert, nil
+}
+
+func RetrieveServiceCABundle(client *kubernetes.Clientset, configMapName, nsName string) ([]byte, error) {
+	// Prompt the injection of the ca bundle into a configmap
+	err := createAnnotatedCABundleInjectionConfigMap(client, configMapName, nsName)
+	if err != nil {
+		return nil, fmt.Errorf("error creating annotated configmap: %v", err)
+	}
+
+	// Retrieve the ca bundle
+	expectedDataSize := 1
+	var ca []byte
+	err = wait.PollImmediate(time.Second*1, time.Second*30, func() (bool, error) {
+		configMap, err := client.CoreV1().ConfigMaps(nsName).Get(configMapName, metav1.GetOptions{})
+		if err != nil {
+			return false, err
+		}
+		if len(configMap.Data) != expectedDataSize {
+			return false, fmt.Errorf("expected data size %d, got %d", expectedDataSize, len(configMap.Data))
+		}
+		_, ok := configMap.Data[serviceCAinjectionDataKey]
+		if !ok {
+			return false, fmt.Errorf("key %q is missing", serviceCAinjectionDataKey)
+		}
+		ca = []byte(configMap.Data[serviceCAinjectionDataKey])
+		return true, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return ca, nil
+}
+
+func createAnnotatedCABundleInjectionConfigMap(client *kubernetes.Clientset, configMapName, nsName string) error {
+	_, err := client.CoreV1().ConfigMaps(nsName).Create(&corev1.ConfigMap{
+		TypeMeta: metav1.TypeMeta{},
+		ObjectMeta: metav1.ObjectMeta{
+			Name: configMapName,
+			Annotations: map[string]string{
+				serviceCAinjectCABundleAnnotationName: "true",
+			},
+		},
+	})
+	return err
+}
+
+func createServingCertAnnotatedService(client *kubernetes.Clientset, secretName, serviceName, namespace string) error {
+	_, err := client.CoreV1().Services(namespace).Create(&corev1.Service{
+		TypeMeta: metav1.TypeMeta{},
+		ObjectMeta: metav1.ObjectMeta{
+			Name: serviceName,
+			Annotations: map[string]string{
+				"service.beta.openshift.io/serving-cert-secret-name": secretName,
+			},
+		},
+		Spec: corev1.ServiceSpec{
+			Ports: []corev1.ServicePort{
+				{
+					Name: "tests",
+					Port: 8443,
+				},
+			},
+		},
+	})
+	return err
+}
+
+func pollForServiceServingSecret(client *kubernetes.Clientset, secretName, namespace string) error {
+	return wait.PollImmediate(time.Second, 10*time.Second, func() (bool, error) {
+		_, err := client.CoreV1().Secrets(namespace).Get(secretName, metav1.GetOptions{})
+		if err != nil && kerrs.IsNotFound(err) {
+			return false, nil
+		}
+		if err != nil {
+			return false, err
+		}
+		return true, nil
+	})
+}
+
+func getServiceServingCertSecretData(client *kubernetes.Clientset, secretName, namespace string) (tls.Certificate, error) {
+	sss, err := client.CoreV1().Secrets(namespace).Get(secretName, metav1.GetOptions{})
+	if err != nil {
+		return tls.Certificate{}, err
+	}
+	keyFile, err := ioutil.TempFile(os.TempDir(), "test-keyfile")
+	if err != nil {
+		return tls.Certificate{}, err
+	}
+	certFile, err := ioutil.TempFile(os.TempDir(), "test-certfile")
+	if err != nil {
+		return tls.Certificate{}, err
+	}
+
+	defer os.Remove(keyFile.Name())
+	defer os.Remove(certFile.Name())
+
+	if _, err = keyFile.Write(sss.Data[corev1.TLSPrivateKeyKey]); err != nil {
+		return tls.Certificate{}, err
+	}
+	if _, err = certFile.Write(sss.Data[corev1.TLSCertKey]); err != nil {
+		return tls.Certificate{}, err
+	}
+
+	if err := keyFile.Close(); err != nil {
+		return tls.Certificate{}, err
+	}
+	if err := certFile.Close(); err != nil {
+		return tls.Certificate{}, err
+	}
+
+	cert, err := tls.LoadX509KeyPair(certFile.Name(), keyFile.Name())
+	if err != nil {
+		return tls.Certificate{}, err
+	}
+	return cert, nil
+}