From bf70e5bd3be297a552dc4b1a68f7fe2d069a57b1 Mon Sep 17 00:00:00 2001
From: Kohei YOSHIDA <kohei@yosida95.com>
Date: Fri, 10 Feb 2023 14:39:35 +0900
Subject: [PATCH] implement UNIX domain socket support for multi-target
 scraping

Signed-off-by: Kohei YOSHIDA <kohei@yosida95.com>
---
 config/config.go      | 18 +++++++++++++-----
 config/config_test.go | 16 ++++++++++++----
 mysqld_exporter.go    |  2 +-
 probe.go              |  6 +++++-
 4 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/config/config.go b/config/config.go
index 5ea55912c..918c98236 100644
--- a/config/config.go
+++ b/config/config.go
@@ -170,12 +170,12 @@ func (m MySqlConfig) validateConfig() error {
 	return nil
 }
 
-func (m MySqlConfig) FormDSN(target string) (string, error) {
+func (m MySqlConfig) FormDSN(proto, target string) (string, error) {
 	config := mysql.NewConfig()
 	config.User = m.User
 	config.Passwd = m.Password
 	config.Net = "tcp"
-	if target == "" {
+	if proto == "" || target == "" {
 		if m.Socket == "" {
 			host := "127.0.0.1"
 			if m.Host != "" {
@@ -191,10 +191,18 @@ func (m MySqlConfig) FormDSN(target string) (string, error) {
 			config.Addr = m.Socket
 		}
 	} else {
-		if _, _, err = net.SplitHostPort(target); err != nil {
-			return "", fmt.Errorf("failed to parse target: %s", err)
+		switch proto {
+		case "tcp":
+			if _, _, err = net.SplitHostPort(target); err != nil {
+				return "", fmt.Errorf("failed to parse target: %s", err)
+			}
+			config.Addr = target
+		case "unix":
+			config.Net = proto
+			config.Addr = target
+		default:
+			return "", fmt.Errorf("unrecognized target protocol: %q", proto)
 		}
-		config.Addr = target
 	}
 
 	if m.SslCa != "" {
diff --git a/config/config_test.go b/config/config_test.go
index b7b1280b2..e0ac6f817 100644
--- a/config/config_test.go
+++ b/config/config_test.go
@@ -154,19 +154,27 @@ func TestFormDSN(t *testing.T) {
 		}
 		convey.Convey("Default Client", func() {
 			cfg := c.GetConfig()
-			section, _ := cfg.Sections["client"]
-			if dsn, err = section.FormDSN(""); err != nil {
+			section := cfg.Sections["client"]
+			if dsn, err = section.FormDSN("", ""); err != nil {
 				t.Error(err)
 			}
 			convey.So(dsn, convey.ShouldEqual, "root:abc@tcp(server2:3306)/")
 		})
 		convey.Convey("Target specific with explicit port", func() {
 			cfg := c.GetConfig()
-			section, _ := cfg.Sections["client.server1"]
-			if dsn, err = section.FormDSN("server1:5000"); err != nil {
+			section := cfg.Sections["client.server1"]
+			if dsn, err = section.FormDSN("tcp", "server1:5000"); err != nil {
 				t.Error(err)
 			}
 			convey.So(dsn, convey.ShouldEqual, "test:foo@tcp(server1:5000)/")
 		})
+		convey.Convey("UNIX domain socket", func() {
+			cfg := c.GetConfig()
+			section := cfg.Sections["client.server1"]
+			if dsn, err = section.FormDSN("unix", "/run/mysqld/mysqld.sock"); err != nil {
+				t.Error(err)
+			}
+			convey.So(dsn, convey.ShouldEqual, "test:foo@unix(/run/mysqld/mysqld.sock)/")
+		})
 	})
 }
diff --git a/mysqld_exporter.go b/mysqld_exporter.go
index c511a8e70..846998cf7 100644
--- a/mysqld_exporter.go
+++ b/mysqld_exporter.go
@@ -141,7 +141,7 @@ func newHandler(metrics collector.Metrics, scrapers []collector.Scraper, logger
 		if !ok {
 			level.Error(logger).Log("msg", "Failed to parse section [client] from config file", "err", err)
 		}
-		if dsn, err = cfgsection.FormDSN(""); err != nil {
+		if dsn, err = cfgsection.FormDSN("", ""); err != nil {
 			level.Error(logger).Log("msg", "Failed to form dsn from section [client]", "err", err)
 		}
 
diff --git a/probe.go b/probe.go
index 47990216c..7f8698f1e 100644
--- a/probe.go
+++ b/probe.go
@@ -36,6 +36,10 @@ func handleProbe(metrics collector.Metrics, scrapers []collector.Scraper, logger
 			http.Error(w, "target is required", http.StatusBadRequest)
 			return
 		}
+		proto := params.Get("protocol")
+		if proto == "" {
+			proto = "tcp"
+		}
 		collectParams := r.URL.Query()["collect[]"]
 
 		if authModule = params.Get("auth_module"); authModule == "" {
@@ -48,7 +52,7 @@ func handleProbe(metrics collector.Metrics, scrapers []collector.Scraper, logger
 			level.Error(logger).Log("msg", fmt.Sprintf("Failed to parse section [%s] from config file", authModule), "err", err)
 			http.Error(w, fmt.Sprintf("Error parsing config section [%s]", authModule), http.StatusBadRequest)
 		}
-		if dsn, err = cfgsection.FormDSN(target); err != nil {
+		if dsn, err = cfgsection.FormDSN(proto, target); err != nil {
 			level.Error(logger).Log("msg", fmt.Sprintf("Failed to form dsn from section [%s]", authModule), "err", err)
 			http.Error(w, fmt.Sprintf("Error forming dsn from config section [%s]", authModule), http.StatusBadRequest)
 		}