Improve PostgreSQL version checker functions

Author: Diogo Nicoleti

Diff for: gauges/backends.go

@@ -123,7 +123,7 @@ type backendsByWaitEventType struct {
 }
 
 func (g *Gauges) backendsByWaitEventTypeQuery() string {
-	if postgres.Version(g.version()).Is96Or10() {
+	if postgres.Version(g.version()).IsEqualOrGreaterThan96() {
 		return `
 			SELECT
 			  COUNT(*) AS total,

Diff for: gauges/gauge.go

@@ -149,11 +149,11 @@ func (g *Gauges) queryWithTimeout(
 	return err
 }
 
-func (g *Gauges) version() string {
-	var version string
-	if err := g.db.QueryRow("show server_version").Scan(&version); err != nil {
+func (g *Gauges) version() int {
+	var version int
+	if err := g.db.QueryRow("show server_version_num").Scan(&version); err != nil {
 		log.WithField("db", g.name).WithError(err).Error("failed to get postgresql version, assuming 9.6.0")
-		return "9.6.0"
+		return 90600
 	}
 	return version
 }

Diff for: gauges/vacuum.go

@@ -112,7 +112,7 @@ func (g *Gauges) VacuumRunningTotal() prometheus.Gauge {
 		SELECT COUNT(*) FROM pg_stat_progress_vacuum WHERE datname = current_database()
 	`
 
-	if !postgres.Version(g.version()).Is96Or10() {
+	if !postgres.Version(g.version()).IsEqualOrGreaterThan96() {
 		log.WithField("db", g.name).
 			Warn("postgresql_vacuum_running_total disabled because it's only supported for PostgreSQL 9.6 or newer versions")
 		return prometheus.NewGauge(gaugeOpts)

Diff for: postgres/version.go

@@ -1,14 +1,22 @@
 package postgres
 
-import "strings"
-
 // Version for a postgres server
-type Version string
+type Version int
+
+// IsEqualOrGreaterThan96 returns whether this is version is greater than 9.6.x
+func (v Version) IsEqualOrGreaterThan96() bool {
+	return v >= 90600
+}
+
+// IsEqualOrGreaterThan10 returns whether this is version is greater than 10.x
+func (v Version) IsEqualOrGreaterThan10() bool {
+	return v >= 100000
+}
 
 // IsWalReplayPausedFunctionName returns the name of the function to verify whether the replication
 // log is paused according to the postgres version
 func (v Version) IsWalReplayPausedFunctionName() string {
-	if v.is10() {
+	if v.IsEqualOrGreaterThan10() {
 		return "pg_is_wal_replay_paused"
 	}
 	return "pg_is_xlog_replay_paused"
@@ -17,7 +25,7 @@ func (v Version) IsWalReplayPausedFunctionName() string {
 // LastWalReceivedLsnFunctionName returns the name of the function that returns the last write-ahead
 // log location received and synced to disk by replication according to the postgres version
 func (v Version) LastWalReceivedLsnFunctionName() string {
-	if v.is10() {
+	if v.IsEqualOrGreaterThan10() {
 		return "pg_last_wal_receive_lsn"
 	}
 	return "pg_last_xlog_receive_location"
@@ -26,7 +34,7 @@ func (v Version) LastWalReceivedLsnFunctionName() string {
 // WalLsnDiffFunctionName returns the name of the function that returns the difference between two write-ahead
 // log locations
 func (v Version) WalLsnDiffFunctionName() string {
-	if v.is10() {
+	if v.IsEqualOrGreaterThan10() {
 		return "pg_wal_lsn_diff"
 	}
 	return "pg_xlog_location_diff"
@@ -35,22 +43,9 @@ func (v Version) WalLsnDiffFunctionName() string {
 // LastWalReplayedLsnFunctionName returns the name of the function that returns the last write-ahead
 // log location replayed during recovery according to the postgres version
 func (v Version) LastWalReplayedLsnFunctionName() string {
-	if v.is10() {
+	if v.IsEqualOrGreaterThan10() {
 		return "pg_last_wal_replay_lsn"
 	}
 	return "pg_last_xlog_replay_location"
 
 }
-
-func (v Version) is96() bool {
-	return strings.HasPrefix(string(v), "9.6.")
-}
-
-func (v Version) is10() bool {
-	return strings.HasPrefix(string(v), "10.")
-}
-
-// Is96Or10 returns whether this is version 9.6.x, 10.x or not
-func (v Version) Is96Or10() bool {
-	return v.is96() || v.is10()
-}

Diff for: postgres/version_test.go

@@ -7,30 +7,46 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestFunctionVersions(t *testing.T) {
-	is96 := Version.is96
-	is10 := Version.is10
-	is96Or10 := Version.Is96Or10
+func TestIsEqualOrGreaterThan96(t *testing.T) {
+	tt := []struct {
+		version  int
+		expected bool
+	}{
+		{90407, false},
+		{90600, true},
+		{90606, true},
+		{100000, true},
+		{100004, true},
+		{110000, true},
+		{110004, true},
+	}
+
+	for _, tc := range tt {
+		testName := fmt.Sprintf("expecting IsEqualOrGreaterThan96(\"%v\") to be %v", tc.version, tc.expected)
+		t.Run(testName, func(t *testing.T) {
+			assert.Equal(t, Version(tc.version).IsEqualOrGreaterThan96(), tc.expected)
+		})
+	}
+}
 
+func TestIsEqualOrGreaterThan10(t *testing.T) {
 	tt := []struct {
-		str      string
-		fn       func(Version) bool
+		version  int
 		expected bool
 	}{
-		{"9.6.6", is96, true},
-		{"9.5.4", is96, false},
-		{"10.3", is10, true},
-		{"9.6.6", is10, false},
-		{"9.6.6", is96Or10, true},
-		{"9.5.4", is96Or10, false},
-		{"10.4", is96Or10, true},
-		{"11.0", is96Or10, false},
+		{90407, false},
+		{90600, false},
+		{90606, false},
+		{100000, true},
+		{100004, true},
+		{110000, true},
+		{110004, true},
 	}
 
 	for _, tc := range tt {
-		testName := fmt.Sprintf("expecting %p(\"%v\") to be %v", tc.fn, tc.str, tc.expected)
+		testName := fmt.Sprintf("expecting IsEqualOrGreaterThan10(\"%v\") to be %v", tc.version, tc.expected)
 		t.Run(testName, func(t *testing.T) {
-			assert.Equal(t, tc.fn(Version(tc.str)), tc.expected)
+			assert.Equal(t, Version(tc.version).IsEqualOrGreaterThan10(), tc.expected)
 		})
 	}
 }