From d07635c5040343906f2a96a51d7dd67dab606104 Mon Sep 17 00:00:00 2001
From: Tamal Saha <tamal@appscode.com>
Date: Fri, 19 May 2017 18:35:30 -0700
Subject: [PATCH] Refactor exporter so that it can be used as a library.

---
 Makefile                                      |  2 +-
 pg_setting.go => collector/pg_setting.go      |  2 +-
 .../pg_setting_test.go                        |  2 +-
 .../postgres_exporter.go                      | 77 +++----------------
 .../postgres_exporter_integration_test.go     | 16 ++--
 .../postgres_exporter_test.go                 | 14 ++--
 main.go                                       | 66 ++++++++++++++++
 7 files changed, 95 insertions(+), 84 deletions(-)
 rename pg_setting.go => collector/pg_setting.go (99%)
 rename pg_setting_test.go => collector/pg_setting_test.go (99%)
 rename postgres_exporter.go => collector/postgres_exporter.go (95%)
 rename postgres_exporter_integration_test.go => collector/postgres_exporter_integration_test.go (95%)
 rename postgres_exporter_test.go => collector/postgres_exporter_test.go (94%)
 create mode 100644 main.go

diff --git a/Makefile b/Makefile
index e2b07efc0..628413ea0 100644
--- a/Makefile
+++ b/Makefile
@@ -89,7 +89,7 @@ fmt: tools
 postgres_exporter_integration_test: $(GO_SRC)
 	CGO_ENABLED=0 go test -c -tags integration \
 		-a -ldflags "-extldflags '-static' -X main.Version=$(VERSION)" \
-		-o postgres_exporter_integration_test -cover -covermode count .
+		-o postgres_exporter_integration_test -cover -covermode count ./collector/...
 
 test: tools
 	@mkdir -p $(COVERDIR)
diff --git a/pg_setting.go b/collector/pg_setting.go
similarity index 99%
rename from pg_setting.go
rename to collector/pg_setting.go
index 0f770d660..a20166b29 100644
--- a/pg_setting.go
+++ b/collector/pg_setting.go
@@ -1,4 +1,4 @@
-package main
+package collector
 
 import (
 	"database/sql"
diff --git a/pg_setting_test.go b/collector/pg_setting_test.go
similarity index 99%
rename from pg_setting_test.go
rename to collector/pg_setting_test.go
index 8b2019edc..ede3765e0 100644
--- a/pg_setting_test.go
+++ b/collector/pg_setting_test.go
@@ -1,6 +1,6 @@
 // +build !integration
 
-package main
+package collector
 
 import (
 	dto "github.com/prometheus/client_model/go"
diff --git a/postgres_exporter.go b/collector/postgres_exporter.go
similarity index 95%
rename from postgres_exporter.go
rename to collector/postgres_exporter.go
index 8a8395c40..6192a7b1b 100644
--- a/postgres_exporter.go
+++ b/collector/postgres_exporter.go
@@ -1,41 +1,25 @@
-package main
+package collector
 
 import (
+	"crypto/sha256"
 	"database/sql"
 	"errors"
 	"fmt"
 	"io/ioutil"
 	"math"
-	"net/http"
 	"net/url"
 	"os"
 	"regexp"
-	"runtime"
 	"strconv"
 	"strings"
 	"sync"
 	"time"
 
-	"gopkg.in/alecthomas/kingpin.v2"
-	"gopkg.in/yaml.v2"
-
-	"crypto/sha256"
 	"github.com/blang/semver"
 	_ "github.com/lib/pq"
 	"github.com/prometheus/client_golang/prometheus"
-	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/prometheus/common/log"
-)
-
-// Version is set during build to the git describe version
-// (semantic version)-(commitish) form.
-var Version = "0.0.1"
-
-var (
-	listenAddress = kingpin.Flag("web.listen-address", "Address to listen on for web interface and telemetry.").Default(":9187").OverrideDefaultFromEnvar("PG_EXPORTER_WEB_LISTEN_ADDRESS").String()
-	metricPath    = kingpin.Flag("web.telemetry-path", "Path under which to expose metrics.").Default("/metrics").OverrideDefaultFromEnvar("PG_EXPORTER_WEB_TELEMETRY_PATH").String()
-	queriesPath   = kingpin.Flag("extend.query-path", "Path to custom queries to run.").Default("").OverrideDefaultFromEnvar("PG_EXPORTER_EXTEND_QUERY_PATH").String()
-	onlyDumpMaps  = kingpin.Flag("dumpmaps", "Do not run, simply dump the maps.").Bool()
+	"gopkg.in/yaml.v2"
 )
 
 // Metric name parts.
@@ -124,7 +108,7 @@ type MetricMap struct {
 }
 
 // TODO: revisit this with the semver system
-func dumpMaps() {
+func DumpMaps() {
 	// TODO: make this function part of the exporter
 	for name, cmap := range builtinMetricMaps {
 		query, ok := queryOverrides[name]
@@ -727,6 +711,12 @@ func NewExporter(dsn string, userQueriesPath string) *Exporter {
 	}
 }
 
+func (e *Exporter) Close() {
+	if e.dbConnection != nil {
+		e.dbConnection.Close() // nolint: errcheck
+	}
+}
+
 // Describe implements prometheus.Collector.
 func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
 	// We cannot know in advance what metrics the exporter will generate
@@ -1038,7 +1028,7 @@ func (e *Exporter) scrape(ch chan<- prometheus.Metric) {
 // DATA_SOURCE_NAME always wins so we do not break older versions
 // reading secrets from files wins over secrets in environment variables
 // DATA_SOURCE_NAME > DATA_SOURCE_{USER|FILE}_FILE > DATA_SOURCE_{USER|FILE}
-func getDataSource() string {
+func GetDataSource() string {
 	var dsn = os.Getenv("DATA_SOURCE_NAME")
 	if len(dsn) == 0 {
 		var user string
@@ -1070,48 +1060,3 @@ func getDataSource() string {
 
 	return dsn
 }
-
-func main() {
-	kingpin.Version(fmt.Sprintf("postgres_exporter %s (built with %s)\n", Version, runtime.Version()))
-	log.AddFlags(kingpin.CommandLine)
-	kingpin.Parse()
-
-	// landingPage contains the HTML served at '/'.
-	// TODO: Make this nicer and more informative.
-	var landingPage = []byte(`<html>
-	<head><title>Postgres exporter</title></head>
-	<body>
-	<h1>Postgres exporter</h1>
-	<p><a href='http://github.com/%60%20+%20*metricPath%20+%20%60'>Metrics</a></p>
-	</body>
-	</html>
-	`)
-
-	if *onlyDumpMaps {
-		dumpMaps()
-		return
-	}
-
-	dsn := getDataSource()
-	if len(dsn) == 0 {
-		log.Fatal("couldn't find environment variables describing the datasource to use")
-	}
-
-	exporter := NewExporter(dsn, *queriesPath)
-	defer func() {
-		if exporter.dbConnection != nil {
-			exporter.dbConnection.Close() // nolint: errcheck
-		}
-	}()
-
-	prometheus.MustRegister(exporter)
-
-	http.Handle(*metricPath, promhttp.Handler())
-	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Content-Type", "Content-Type:text/plain; charset=UTF-8") // nolint: errcheck
-		w.Write(landingPage)                                                     // nolint: errcheck
-	})
-
-	log.Infof("Starting Server: %s", *listenAddress)
-	log.Fatal(http.ListenAndServe(*listenAddress, nil))
-}
diff --git a/postgres_exporter_integration_test.go b/collector/postgres_exporter_integration_test.go
similarity index 95%
rename from postgres_exporter_integration_test.go
rename to collector/postgres_exporter_integration_test.go
index 3387f2b56..bbe014103 100644
--- a/postgres_exporter_integration_test.go
+++ b/collector/postgres_exporter_integration_test.go
@@ -3,19 +3,17 @@
 // working.
 // +build integration
 
-package main
+package collector
 
 import (
-	"os"
-	"testing"
-
-	. "gopkg.in/check.v1"
-
 	"database/sql"
 	"fmt"
+	"os"
+	"testing"
 
 	_ "github.com/lib/pq"
 	"github.com/prometheus/client_golang/prometheus"
+	. "gopkg.in/check.v1"
 )
 
 // Hook up gocheck into the "go test" runner.
@@ -78,6 +76,8 @@ func (s *IntegrationSuite) TestAllNamespacesReturnResults(c *C) {
 // the exporter. Related to https://github.com/wrouesnel/postgres_exporter/issues/93
 // although not a replication of the scenario.
 func (s *IntegrationSuite) TestInvalidDsnDoesntCrash(c *C) {
+	queriesPath := os.Getenv("PG_EXPORTER_EXTEND_QUERY_PATH")
+
 	// Setup a dummy channel to consume metrics
 	ch := make(chan prometheus.Metric, 100)
 	go func() {
@@ -86,12 +86,12 @@ func (s *IntegrationSuite) TestInvalidDsnDoesntCrash(c *C) {
 	}()
 
 	// Send a bad DSN
-	exporter := NewExporter("invalid dsn", *queriesPath)
+	exporter := NewExporter("invalid dsn", queriesPath)
 	c.Assert(exporter, NotNil)
 	exporter.scrape(ch)
 
 	// Send a DSN to a non-listening port.
-	exporter = NewExporter("postgresql://nothing:nothing@127.0.0.1:1/nothing", *queriesPath)
+	exporter = NewExporter("postgresql://nothing:nothing@127.0.0.1:1/nothing", queriesPath)
 	c.Assert(exporter, NotNil)
 	exporter.scrape(ch)
 }
diff --git a/postgres_exporter_test.go b/collector/postgres_exporter_test.go
similarity index 94%
rename from postgres_exporter_test.go
rename to collector/postgres_exporter_test.go
index e5a50b098..94c3dd7d9 100644
--- a/postgres_exporter_test.go
+++ b/collector/postgres_exporter_test.go
@@ -1,6 +1,6 @@
 // +build !integration
 
-package main
+package collector
 
 import (
 	. "gopkg.in/check.v1"
@@ -89,11 +89,11 @@ func (s *FunctionalSuite) TestSemanticVersionColumnDiscard(c *C) {
 // test read username and password from file
 func (s *FunctionalSuite) TestEnvironmentSettingWithSecretsFiles(c *C) {
 
-	err := os.Setenv("DATA_SOURCE_USER_FILE", "./tests/username_file")
+	err := os.Setenv("DATA_SOURCE_USER_FILE", "../tests/username_file")
 	c.Assert(err, IsNil)
 	defer UnsetEnvironment(c, "DATA_SOURCE_USER_FILE")
 
-	err = os.Setenv("DATA_SOURCE_PASS_FILE", "./tests/userpass_file")
+	err = os.Setenv("DATA_SOURCE_PASS_FILE", "../tests/userpass_file")
 	c.Assert(err, IsNil)
 	defer UnsetEnvironment(c, "DATA_SOURCE_PASS_FILE")
 
@@ -103,7 +103,7 @@ func (s *FunctionalSuite) TestEnvironmentSettingWithSecretsFiles(c *C) {
 
 	var expected = "postgresql://custom_username:custom_password@localhost:5432/?sslmode=disable"
 
-	dsn := getDataSource()
+	dsn := GetDataSource()
 	if dsn != expected {
 		c.Errorf("Expected Username to be read from file. Found=%v, expected=%v", dsn, expected)
 	}
@@ -117,7 +117,7 @@ func (s *FunctionalSuite) TestEnvironmentSettingWithDns(c *C) {
 	c.Assert(err, IsNil)
 	defer UnsetEnvironment(c, "DATA_SOURCE_NAME")
 
-	dsn := getDataSource()
+	dsn := GetDataSource()
 	if dsn != envDsn {
 		c.Errorf("Expected Username to be read from file. Found=%v, expected=%v", dsn, envDsn)
 	}
@@ -131,7 +131,7 @@ func (s *FunctionalSuite) TestEnvironmentSettingWithDnsAndSecrets(c *C) {
 	c.Assert(err, IsNil)
 	defer UnsetEnvironment(c, "DATA_SOURCE_NAME")
 
-	err = os.Setenv("DATA_SOURCE_USER_FILE", "./tests/username_file")
+	err = os.Setenv("DATA_SOURCE_USER_FILE", "../tests/username_file")
 	c.Assert(err, IsNil)
 	defer UnsetEnvironment(c, "DATA_SOURCE_USER_FILE")
 
@@ -139,7 +139,7 @@ func (s *FunctionalSuite) TestEnvironmentSettingWithDnsAndSecrets(c *C) {
 	c.Assert(err, IsNil)
 	defer UnsetEnvironment(c, "DATA_SOURCE_PASS")
 
-	dsn := getDataSource()
+	dsn := GetDataSource()
 	if dsn != envDsn {
 		c.Errorf("Expected Username to be read from file. Found=%v, expected=%v", dsn, envDsn)
 	}
diff --git a/main.go b/main.go
new file mode 100644
index 000000000..1b11b9dde
--- /dev/null
+++ b/main.go
@@ -0,0 +1,66 @@
+package main
+
+import (
+	"fmt"
+	"net/http"
+	"runtime"
+
+	_ "github.com/lib/pq"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"github.com/prometheus/common/log"
+	"github.com/wrouesnel/postgres_exporter/collector"
+	"gopkg.in/alecthomas/kingpin.v2"
+)
+
+// Version is set during build to the git describe version
+// (semantic version)-(commitish) form.
+var Version = "0.0.1"
+
+var (
+	listenAddress = kingpin.Flag("web.listen-address", "Address to listen on for web interface and telemetry.").Default(":9187").OverrideDefaultFromEnvar("PG_EXPORTER_WEB_LISTEN_ADDRESS").String()
+	metricPath    = kingpin.Flag("web.telemetry-path", "Path under which to expose metrics.").Default("/metrics").OverrideDefaultFromEnvar("PG_EXPORTER_WEB_TELEMETRY_PATH").String()
+	queriesPath   = kingpin.Flag("extend.query-path", "Path to custom queries to run.").Default("").OverrideDefaultFromEnvar("PG_EXPORTER_EXTEND_QUERY_PATH").String()
+	onlyDumpMaps  = kingpin.Flag("dumpmaps", "Do not run, simply dump the maps.").Bool()
+)
+
+func main() {
+	kingpin.Version(fmt.Sprintf("postgres_exporter %s (built with %s)\n", Version, runtime.Version()))
+	log.AddFlags(kingpin.CommandLine)
+	kingpin.Parse()
+
+	// landingPage contains the HTML served at '/'.
+	// TODO: Make this nicer and more informative.
+	var landingPage = []byte(`<html>
+	<head><title>Postgres exporter</title></head>
+	<body>
+	<h1>Postgres exporter</h1>
+	<p><a href='http://github.com/%60%20+%20*metricPath%20+%20%60'>Metrics</a></p>
+	</body>
+	</html>
+	`)
+
+	if *onlyDumpMaps {
+		collector.DumpMaps()
+		return
+	}
+
+	dsn := collector.GetDataSource()
+	if len(dsn) == 0 {
+		log.Fatal("couldn't find environment variables describing the datasource to use")
+	}
+
+	exporter := collector.NewExporter(dsn, *queriesPath)
+	defer exporter.Close()
+
+	prometheus.MustRegister(exporter)
+
+	http.Handle(*metricPath, promhttp.Handler())
+	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "Content-Type:text/plain; charset=UTF-8") // nolint: errcheck
+		w.Write(landingPage)                                                     // nolint: errcheck
+	})
+
+	log.Infof("Starting Server: %s", *listenAddress)
+	log.Fatal(http.ListenAndServe(*listenAddress, nil))
+}