Skip to content

Commit d983d32

Browse files
SuperQSuperQ
authored
Merge pull request #785 from sysadmind/collector-instance
Add the instance struct to handle connections
2 parents 5db7cfb + ab33346 commit d983d32

19 files changed

+139
-44
lines changed

‎collector/collector.go‎

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ package collector
1515

1616
import (
1717
"context"
18-
"database/sql"
1918
"errors"
2019
"fmt"
2120
"sync"
@@ -59,7 +58,7 @@ var (
5958
)
6059

6160
type Collector interface {
62-
Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error
61+
Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error
6362
}
6463

6564
type collectorConfig struct {
@@ -92,7 +91,7 @@ type PostgresCollector struct {
9291
Collectors map[string]Collector
9392
logger log.Logger
9493

95-
db *sql.DB
94+
instance *instance
9695
}
9796

9897
type Option func(*PostgresCollector) error
@@ -149,14 +148,11 @@ func NewPostgresCollector(logger log.Logger, excludeDatabases []string, dsn stri
149148
return nil, errors.New("empty dsn")
150149
}
151150

152-
db, err := sql.Open("postgres", dsn)
151+
instance, err := newInstance(dsn)
153152
if err != nil {
154153
return nil, err
155154
}
156-
db.SetMaxOpenConns(1)
157-
db.SetMaxIdleConns(1)
158-
159-
p.db = db
155+
p.instance = instance
160156

161157
return p, nil
162158
}
@@ -174,16 +170,16 @@ func (p PostgresCollector) Collect(ch chan<- prometheus.Metric) {
174170
wg.Add(len(p.Collectors))
175171
for name, c := range p.Collectors {
176172
go func(name string, c Collector) {
177-
execute(ctx, name, c, p.db, ch, p.logger)
173+
execute(ctx, name, c, p.instance, ch, p.logger)
178174
wg.Done()
179175
}(name, c)
180176
}
181177
wg.Wait()
182178
}
183179

184-
func execute(ctx context.Context, name string, c Collector, db *sql.DB, ch chan<- prometheus.Metric, logger log.Logger) {
180+
func execute(ctx context.Context, name string, c Collector, instance *instance, ch chan<- prometheus.Metric, logger log.Logger) {
185181
begin := time.Now()
186-
err := c.Update(ctx, db, ch)
182+
err := c.Update(ctx, instance, ch)
187183
duration := time.Since(begin)
188184
var success float64
189185

‎collector/instance.go‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
// Copyright 2023 The Prometheus Authors
2+
// Licensed under the Apache License, Version 2.0 (the "License");
3+
// you may not use this file except in compliance with the License.
4+
// You may obtain a copy of the License at
5+
//
6+
// http://www.apache.org/licenses/LICENSE-2.0
7+
//
8+
// Unless required by applicable law or agreed to in writing, software
9+
// distributed under the License is distributed on an "AS IS" BASIS,
10+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
// See the License for the specific language governing permissions and
12+
// limitations under the License.
13+
14+
package collector
15+
16+
import (
17+
"database/sql"
18+
"fmt"
19+
"regexp"
20+
21+
"github.com/blang/semver/v4"
22+
)
23+
24+
type instance struct {
25+
db *sql.DB
26+
version semver.Version
27+
}
28+
29+
func newInstance(dsn string) (*instance, error) {
30+
i := &instance{}
31+
db, err := sql.Open("postgres", dsn)
32+
if err != nil {
33+
return nil, err
34+
}
35+
db.SetMaxOpenConns(1)
36+
db.SetMaxIdleConns(1)
37+
i.db = db
38+
39+
version, err := queryVersion(db)
40+
if err != nil {
41+
db.Close()
42+
return nil, err
43+
}
44+
45+
i.version = version
46+
47+
return i, nil
48+
}
49+
50+
func (i *instance) getDB() *sql.DB {
51+
return i.db
52+
}
53+
54+
func (i *instance) Close() error {
55+
return i.db.Close()
56+
}
57+
58+
// Regex used to get the "short-version" from the postgres version field.
59+
// The result of SELECT version() is something like "PostgreSQL 9.6.2 on x86_64-pc-linux-gnu, compiled by gcc (GCC) 6.2.1 20160830, 64-bit"
60+
var versionRegex = regexp.MustCompile(`^\w+ ((\d+)(\.\d+)?(\.\d+)?)`)
61+
var serverVersionRegex = regexp.MustCompile(`^((\d+)(\.\d+)?(\.\d+)?)`)
62+
63+
func queryVersion(db *sql.DB) (semver.Version, error) {
64+
var version string
65+
err := db.QueryRow("SELECT version();").Scan(&version)
66+
if err != nil {
67+
return semver.Version{}, err
68+
}
69+
submatches := versionRegex.FindStringSubmatch(version)
70+
if len(submatches) > 1 {
71+
return semver.ParseTolerant(submatches[1])
72+
}
73+
74+
// We could also try to parse the version from the server_version field.
75+
// This is of the format 13.3 (Debian 13.3-1.pgdg100+1)
76+
err = db.QueryRow("SHOW server_version;").Scan(&version)
77+
if err != nil {
78+
return semver.Version{}, err
79+
}
80+
submatches = serverVersionRegex.FindStringSubmatch(version)
81+
if len(submatches) > 1 {
82+
return semver.ParseTolerant(submatches[1])
83+
}
84+
return semver.Version{}, fmt.Errorf("could not parse version from %q", version)
85+
}

‎collector/pg_database.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ package collector
1515

1616
import (
1717
"context"
18-
"database/sql"
1918

2019
"github.com/go-kit/log"
2120
"github.com/prometheus/client_golang/prometheus"
@@ -66,7 +65,8 @@ var (
6665
// each database individually. This is because we can't filter the
6766
// list of databases in the query because the list of excluded
6867
// databases is dynamic.
69-
func (c PGDatabaseCollector) Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error {
68+
func (c PGDatabaseCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
69+
db := instance.getDB()
7070
// Query the list of databases
7171
rows, err := db.QueryContext(ctx,
7272
pgDatabaseQuery,

‎collector/pg_database_test.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ func TestPGDatabaseCollector(t *testing.T) {
2929
}
3030
defer db.Close()
3131

32+
inst := &instance{db: db}
33+
3234
mock.ExpectQuery(sanitizeQuery(pgDatabaseQuery)).WillReturnRows(sqlmock.NewRows([]string{"datname"}).
3335
AddRow("postgres"))
3436

@@ -39,7 +41,7 @@ func TestPGDatabaseCollector(t *testing.T) {
3941
go func() {
4042
defer close(ch)
4143
c := PGDatabaseCollector{}
42-
if err := c.Update(context.Background(), db, ch); err != nil {
44+
if err := c.Update(context.Background(), inst, ch); err != nil {
4345
t.Errorf("Error calling PGDatabaseCollector.Update: %s", err)
4446
}
4547
}()

‎collector/pg_postmaster.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ package collector
1515

1616
import (
1717
"context"
18-
"database/sql"
1918

2019
"github.com/prometheus/client_golang/prometheus"
2120
)
@@ -47,7 +46,8 @@ var (
4746
pgPostmasterQuery = "SELECT pg_postmaster_start_time from pg_postmaster_start_time();"
4847
)
4948

50-
func (c *PGPostmasterCollector) Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error {
49+
func (c *PGPostmasterCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
50+
db := instance.getDB()
5151
row := db.QueryRowContext(ctx,
5252
pgPostmasterQuery)
5353

‎collector/pg_postmaster_test.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ func TestPgPostmasterCollector(t *testing.T) {
2929
}
3030
defer db.Close()
3131

32+
inst := &instance{db: db}
33+
3234
mock.ExpectQuery(sanitizeQuery(pgPostmasterQuery)).WillReturnRows(sqlmock.NewRows([]string{"pg_postmaster_start_time"}).
3335
AddRow(1685739904))
3436

@@ -37,7 +39,7 @@ func TestPgPostmasterCollector(t *testing.T) {
3739
defer close(ch)
3840
c := PGPostmasterCollector{}
3941

40-
if err := c.Update(context.Background(), db, ch); err != nil {
42+
if err := c.Update(context.Background(), inst, ch); err != nil {
4143
t.Errorf("Error calling PGPostmasterCollector.Update: %s", err)
4244
}
4345
}()

‎collector/pg_process_idle.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ package collector
1515

1616
import (
1717
"context"
18-
"database/sql"
1918

2019
"github.com/go-kit/log"
2120
"github.com/prometheus/client_golang/prometheus"
@@ -42,7 +41,8 @@ var pgProcessIdleSeconds = prometheus.NewDesc(
4241
prometheus.Labels{},
4342
)
4443

45-
func (PGProcessIdleCollector) Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error {
44+
func (PGProcessIdleCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
45+
db := instance.getDB()
4646
row := db.QueryRowContext(ctx,
4747
`WITH
4848
metrics AS (

‎collector/pg_replication_slot.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ package collector
1515

1616
import (
1717
"context"
18-
"database/sql"
1918

2019
"github.com/go-kit/log"
2120
"github.com/prometheus/client_golang/prometheus"
@@ -73,7 +72,8 @@ var (
7372
pg_replication_slots;`
7473
)
7574

76-
func (PGReplicationSlotCollector) Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error {
75+
func (PGReplicationSlotCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
76+
db := instance.getDB()
7777
rows, err := db.QueryContext(ctx,
7878
pgReplicationSlotQuery)
7979
if err != nil {

‎collector/pg_replication_slot_test.go‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ func TestPgReplicationSlotCollectorActive(t *testing.T) {
2929
}
3030
defer db.Close()
3131

32+
inst := &instance{db: db}
33+
3234
columns := []string{"slot_name", "current_wal_lsn", "confirmed_flush_lsn", "active"}
3335
rows := sqlmock.NewRows(columns).
3436
AddRow("test_slot", 5, 3, true)
@@ -39,7 +41,7 @@ func TestPgReplicationSlotCollectorActive(t *testing.T) {
3941
defer close(ch)
4042
c := PGReplicationSlotCollector{}
4143

42-
if err := c.Update(context.Background(), db, ch); err != nil {
44+
if err := c.Update(context.Background(), inst, ch); err != nil {
4345
t.Errorf("Error calling PGPostmasterCollector.Update: %s", err)
4446
}
4547
}()
@@ -68,6 +70,8 @@ func TestPgReplicationSlotCollectorInActive(t *testing.T) {
6870
}
6971
defer db.Close()
7072

73+
inst := &instance{db: db}
74+
7175
columns := []string{"slot_name", "current_wal_lsn", "confirmed_flush_lsn", "active"}
7276
rows := sqlmock.NewRows(columns).
7377
AddRow("test_slot", 6, 12, false)
@@ -78,7 +82,7 @@ func TestPgReplicationSlotCollectorInActive(t *testing.T) {
7882
defer close(ch)
7983
c := PGReplicationSlotCollector{}
8084

81-
if err := c.Update(context.Background(), db, ch); err != nil {
85+
if err := c.Update(context.Background(), inst, ch); err != nil {
8286
t.Errorf("Error calling PGReplicationSlotCollector.Update: %s", err)
8387
}
8488
}()

‎collector/pg_stat_bgwriter.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ package collector
1515

1616
import (
1717
"context"
18-
"database/sql"
1918
"time"
2019

2120
"github.com/prometheus/client_golang/prometheus"
@@ -117,7 +116,8 @@ var (
117116
FROM pg_stat_bgwriter;`
118117
)
119118

120-
func (PGStatBGWriterCollector) Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error {
119+
func (PGStatBGWriterCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
120+
db := instance.getDB()
121121
row := db.QueryRowContext(ctx,
122122
statBGWriterQuery)
123123

0 commit comments

Comments
 (0)