From 8a4afb4ea2059d518b3a626ea216f2bf259368b8 Mon Sep 17 00:00:00 2001
From: Ulric Qin <ulric.qin@gmail.com>
Date: Fri, 10 Jun 2022 13:10:31 +0800
Subject: [PATCH] support template for prometheus label value

---
 conf/input.prometheus/prometheus.toml |  3 ++
 conf/logs.toml                        |  2 +-
 config/urllabel.go                    | 53 +++++++++++++++++++++++++++
 inputs/prometheus/prometheus.go       | 17 ++++++++-
 4 files changed, 73 insertions(+), 2 deletions(-)
 create mode 100644 config/urllabel.go

diff --git a/conf/input.prometheus/prometheus.toml b/conf/input.prometheus/prometheus.toml
index c7527a52..030173cf 100644
--- a/conf/input.prometheus/prometheus.toml
+++ b/conf/input.prometheus/prometheus.toml
@@ -6,6 +6,9 @@
 #     "http://localhost:9104/metrics"
 # ]
 
+url_label_key = "instance"
+url_label_value = "{{.Host}}"
+
 ## Scrape Services available in Consul Catalog
 [instances.consul]
   enabled = false
diff --git a/conf/logs.toml b/conf/logs.toml
index 593226df..a4d8b721 100644
--- a/conf/logs.toml
+++ b/conf/logs.toml
@@ -2,7 +2,7 @@
 ## key 占位符
 api_key = "ef4ahfbwzwwtlwfpbertgq1i6mq0ab1q"
 ## 是否开启日志采集
-enable = true
+enable = false
 ## 接受日志的server地址
 send_to = "127.0.0.1:17878"
 ## 发送日志的协议 http/tcp
diff --git a/config/urllabel.go b/config/urllabel.go
new file mode 100644
index 00000000..b6fa7e35
--- /dev/null
+++ b/config/urllabel.go
@@ -0,0 +1,53 @@
+package config
+
+import (
+	"bytes"
+	"net/url"
+	"text/template"
+)
+
+type UrlLabel struct {
+	LabelKey      string `toml:"url_label_key"`
+	LabelValue    string `toml:"url_label_value"`
+	LabelValueTpl *template.Template
+}
+
+func (ul *UrlLabel) PrepareUrlTemplate() error {
+	if ul.LabelKey == "" {
+		ul.LabelKey = "instance"
+	}
+
+	if ul.LabelValue != "" {
+		var err error
+		ul.LabelValueTpl, err = template.New("v").Parse(ul.LabelValue)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (ul *UrlLabel) GenerateLabel(u *url.URL) (string, string, error) {
+	if ul.LabelValue == "" {
+		return ul.LabelKey, u.String(), nil
+	}
+
+	dict := map[string]string{
+		"Scheme":   u.Scheme,
+		"Host":     u.Host,
+		"Hostname": u.Hostname(),
+		"Port":     u.Port(),
+		"Path":     u.Path,
+		"Query":    u.RawQuery,
+		"Fragment": u.Fragment,
+	}
+
+	var buffer bytes.Buffer
+	err := ul.LabelValueTpl.Execute(&buffer, dict)
+	if err != nil {
+		return "", "", err
+	}
+
+	return ul.LabelKey, buffer.String(), nil
+}
diff --git a/inputs/prometheus/prometheus.go b/inputs/prometheus/prometheus.go
index 6e222180..32bfd6ce 100644
--- a/inputs/prometheus/prometheus.go
+++ b/inputs/prometheus/prometheus.go
@@ -38,6 +38,8 @@ type Instance struct {
 	IgnoreLabelKeys   []string          `toml:"ignore_label_keys"`
 	Headers           []string          `toml:"headers"`
 
+	config.UrlLabel
+
 	ignoreMetricsFilter   filter.Filter
 	ignoreLabelKeysFilter filter.Filter
 	tls.ClientConfig
@@ -81,6 +83,10 @@ func (ins *Instance) Init() error {
 		}
 	}
 
+	if err := ins.PrepareUrlTemplate(); err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -206,7 +212,16 @@ func (p *Prometheus) gatherUrl(urlwg *sync.WaitGroup, slist *list.SafeList, ins
 
 	ins.setHeaders(req)
 
-	labels := map[string]string{"url": u.String()}
+	labels := map[string]string{}
+
+	urlKey, urlVal, err := ins.GenerateLabel(u)
+	if err != nil {
+		log.Println("E! failed to generate url label value:", err)
+		return
+	}
+
+	labels[urlKey] = urlVal
+
 	for key, val := range ins.Labels {
 		labels[key] = val
 	}