Go script that exposes information about statefulsets and deployments in Kubernetes cluster as Prometheus metrics

labels_to_prometheus_metrics.go

package main

import (
	"log"
	"net/http"
	"os"
	"strings"

	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/tools/cache"
)

// Define the Prometheus metrics
var (
	resourceMetric = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Name: "kubernetes_resource_info",
			Help: "Kubernetes resource information",
		},
		[]string{"namespace", "name", "chart", "revision", "version", "container", "image", "tag"},
	)
)

func main() {
	var config *rest.Config
	var err error

	// Use in-cluster configuration if running in a cluster, fallback to kubeconfig otherwise
	if config, err = rest.InClusterConfig(); err != nil {
		kubeconfig := os.Getenv("KUBECONFIG")
		if kubeconfig == "" {
			kubeconfig = os.Getenv("HOME") + "/.kube/config"
		}
		config, err = clientcmd.BuildConfigFromFlags("", kubeconfig)
		if err != nil {
			log.Fatalf("Failed to load Kubernetes config: %v", err)
		}
	}

	// Create Kubernetes client
	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		log.Fatalf("Failed to create Kubernetes client: %v", err)
	}

	// Create Kubernetes informer factory
	factory := informers.NewSharedInformerFactory(clientset, 0)

	// Create deployment informer
	deploymentInformer := factory.Apps().V1().Deployments().Informer()
	deploymentInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc: func(obj interface{}) {
			deployment := obj.(*appsv1.Deployment)
			updateMetrics(deployment, &deployment.Spec.Template.Spec)
		},
		UpdateFunc: func(oldObj, newObj interface{}) {
			deployment := newObj.(*appsv1.Deployment)
			updateMetrics(deployment, &deployment.Spec.Template.Spec)
		},
		DeleteFunc: func(obj interface{}) {
			deployment := obj.(*appsv1.Deployment)
			deleteMetrics(deployment, &deployment.Spec.Template.Spec)
		},
	})

	// Create statefulset informer
	statefulSetInformer := factory.Apps().V1().StatefulSets().Informer()
	statefulSetInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc: func(obj interface{}) {
			statefulSet := obj.(*appsv1.StatefulSet)
			updateMetrics(statefulSet, &statefulSet.Spec.Template.Spec)
		},
		UpdateFunc: func(oldObj, newObj interface{}) {
			statefulSet := newObj.(*appsv1.StatefulSet)
			updateMetrics(statefulSet, &statefulSet.Spec.Template.Spec)
		},
		DeleteFunc: func(obj interface{}) {
			statefulSet := obj.(*appsv1.StatefulSet)
			deleteMetrics(statefulSet, &statefulSet.Spec.Template.Spec)
		},
	})

	// Register Prometheus metrics
	prometheus.MustRegister(resourceMetric)

	// Create a stopCh for graceful shutdown
	stopCh := make(chan struct{})
	defer close(stopCh)

	// Start informers and wait for cache sync
	factory.Start(stopCh)
	factory.WaitForCacheSync(stopCh)

	// Start serving metrics
	http.Handle("/metrics", promhttp.Handler())
	log.Fatal(http.ListenAndServe(":8080", nil))
}

func updateMetrics(obj metav1.Object, spec *corev1.PodSpec) {
	labels := obj.GetLabels()

	for _, container := range spec.Containers {
		parts := strings.SplitN(container.Image, ":", 2)
		image := parts[0]
		tag := ""
		if len(parts) > 1 {
			tag = parts[1]
		}

		resourceMetric.With(prometheus.Labels{
			"namespace": obj.GetNamespace(),
			"name":      obj.GetName(),
			"chart":     labels["helm.sh/chart"],
			"revision":  labels["deployment.kubernetes.io/revision"],
			"version":   labels["version"],
			"container": container.Name,
			"image":     image,
			"tag":       tag,
		}).Set(1)
	}
}

func deleteMetrics(obj metav1.Object, spec *corev1.PodSpec) {
	labels := obj.GetLabels()

	for _, container := range spec.Containers {
		parts := strings.SplitN(container.Image, ":", 2)
		image := parts[0]
		tag := ""
		if len(parts) > 1 {
			tag = parts[1]
		}

		resourceMetric.Delete(prometheus.Labels{
			"namespace": obj.GetNamespace(),
			"name":      obj.GetName(),
			"chart":     labels["helm.sh/chart"],
			"revision":  labels["deployment.kubernetes.io/revision"],
			"version":   labels["version"],
			"container": container.Name,
			"image":     image,
			"tag":       tag,
		})
	}
}