Skip to content

Commit

Permalink
bug-fixing: PostgreSQL failed to check if the server wasn't turn on t…
Browse files Browse the repository at this point in the history 
…he SSL (#286)

* bug-fixing: PostgreSQL failed to check if the server wasn't turn on the SSL

* using a new function to do defer db.Close() in for-loop
  • Loading branch information
@haoel
haoel authored Feb 10, 2023
1 parent eca4619 commit 7026081
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 53 deletions.
127 changes: 75 additions & 52 deletions probe/client/postgres/postgres.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,8 +62,9 @@ func New(opt conf.Options) (*PostgreSQL, error) {
return nil, fmt.Errorf("TLS Config Error - %v", err)
} else if tls != nil {
tls.InsecureSkipVerify = true
clientOptions = append(clientOptions, pgdriver.WithTLSConfig(tls))
}
// if the tls is nil which means `sslmode=disable`
clientOptions = append(clientOptions, pgdriver.WithTLSConfig(tls))

pg := &PostgreSQL{
Options: opt,
Expand Down Expand Up @@ -97,65 +98,87 @@ func (r *PostgreSQL) checkData() error {
func (r *PostgreSQL) Probe() (bool, string) {

if len(r.Data) > 0 {
for k, v := range r.Data {
log.Debugf("[%s / %s / %s] - Verifying Data - [%s] : [%s]", r.ProbeKind, r.ProbeName, r.ProbeTag, k, v)
//connect to the database
dbName, sqlstr, err := r.getSQL(k)
if err != nil {
return false, fmt.Sprintf("Invalid SQL data - [%s], %v", v, err)
}
clientOptions := append(r.ClientOptions, pgdriver.WithDatabase(dbName))
db := sql.OpenDB(pgdriver.NewConnector(clientOptions...))
if db == nil {
return false, "OpenDB error"
}
// query the data
log.Debugf("[%s / %s / %s] - SQL - [%s]", r.ProbeKind, r.ProbeName, r.ProbeTag, sqlstr)
rows, err := db.Query(sqlstr)
if err != nil {
return false, fmt.Sprintf("Query error - [%s], %v", v, err)
}
if !rows.Next() {
rows.Close()
return false, fmt.Sprintf("No data found for [%s]", k)
}
//check the value is equal to the value in data
var value string
if err := rows.Scan(&value); err != nil {
rows.Close()
return false, err.Error()
}
if value != v {
rows.Close()
return false, fmt.Sprintf("Value not match for [%s] expected [%s] got [%s] ", k, v, value)
}
rows.Close()
db.Close()
log.Debugf("[%s / %s / %s] - Data Verified Successfully! - [%s] : [%s]", r.ProbeKind, r.ProbeName, r.ProbeTag, k, v)
}
} else {
r.ClientOptions = append(r.ClientOptions, pgdriver.WithDatabase("template1"))
db := sql.OpenDB(pgdriver.NewConnector(r.ClientOptions...))
if db == nil {
return false, "OpenDB error"
}
defer db.Close()
return r.ProbeWithDataChecking()
}
return r.ProbeWithPing()
}

if err := db.Ping(); err != nil {
return false, err.Error()
}
// ProbeWithPing do the health check with ping & Select 1;
func (r *PostgreSQL) ProbeWithPing() (bool, string) {
r.ClientOptions = append(r.ClientOptions, pgdriver.WithDatabase("template1"))
db := sql.OpenDB(pgdriver.NewConnector(r.ClientOptions...))
if db == nil {
return false, "OpenDB error"
}
defer db.Close()

// run a SQL to test
row, err := db.Query(`SELECT 1`)
if err != nil {
return false, err.Error()
if err := db.Ping(); err != nil {
return false, err.Error()
}

// run a SQL to test
row, err := db.Query(`SELECT 1`)
if err != nil {
return false, err.Error()
}
row.Close()
return true, "Check PostgreSQL Server Successfully!"
}

// ProbeWithDataChecking do the health check with data checking
func (r *PostgreSQL) ProbeWithDataChecking() (bool, string) {
if len(r.Data) == 0 {
log.Warnf("[%s / %s / %s] - No data found, use ping instead", r.ProbeKind, r.ProbeName, r.ProbeTag)
return r.ProbeWithPing()
}

for k, v := range r.Data {
if ok, msg := r.verifyData(k, v); !ok {
return ok, msg
}
row.Close()
}

return true, "Check PostgreSQL Server Successfully!"
}

func (r *PostgreSQL) verifyData(k, v string) (bool, string) {
log.Debugf("[%s / %s / %s] - Verifying Data - [%s] : [%s]", r.ProbeKind, r.ProbeName, r.ProbeTag, k, v)
//connect to the database
dbName, sqlstr, err := r.getSQL(k)
if err != nil {
return false, fmt.Sprintf("Invalid SQL data - [%s], %v", v, err)
}
clientOptions := append(r.ClientOptions, pgdriver.WithDatabase(dbName))
db := sql.OpenDB(pgdriver.NewConnector(clientOptions...))
if db == nil {
return false, "OpenDB error"
}
defer db.Close()

// query the data
log.Debugf("[%s / %s / %s] - SQL - [%s]", r.ProbeKind, r.ProbeName, r.ProbeTag, sqlstr)
rows, err := db.Query(sqlstr)
if err != nil {
return false, fmt.Sprintf("Query error - [%s], %v", v, err)
}
defer rows.Close()

if !rows.Next() {
return false, fmt.Sprintf("No data found for [%s]", k)
}
//check the value is equal to the value in data
var value string
if err := rows.Scan(&value); err != nil {
return false, err.Error()
}
if value != v {
return false, fmt.Sprintf("Value not match for [%s] expected [%s] got [%s] ", k, v, value)
}

log.Debugf("[%s / %s / %s] - Data Verified Successfully! - [%s] : [%s]", r.ProbeKind, r.ProbeName, r.ProbeTag, k, v)
return true, "Check PostgreSQL Server Successfully!"
}

// getSQL get the SQL statement
// input: database:table:column:key:value
// output: SELECT column FROM database.table WHERE key = value
Expand Down
6 changes: 5 additions & 1 deletion probe/client/postgres/postgres_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ func TestPostgreSQL(t *testing.T) {
assert.Equal(t, conf.Timeout(), pgd.Config().DialTimeout)
assert.Equal(t, conf.Timeout(), pgd.Config().ReadTimeout)
assert.Equal(t, conf.Timeout(), pgd.Config().WriteTimeout)
assert.Nil(t, pgd.Config().TLSConfig.ClientCAs)
assert.Nil(t, pgd.Config().TLSConfig)

monkey.Patch(sql.OpenDB, func(c driver.Connector) *sql.DB {
return &sql.DB{}
Expand All @@ -86,6 +86,10 @@ func TestPostgreSQL(t *testing.T) {
assert.True(t, s)
assert.Contains(t, m, "Successfully")

s, m = pg.ProbeWithDataChecking()
assert.True(t, s)
assert.Contains(t, m, "Successfully")

// TLS config success
var tc *global.TLS
monkey.PatchInstanceMethod(reflect.TypeOf(tc), "Config", func(_ *global.TLS) (*tls.Config, error) {
Expand Down

0 comments on commit 7026081

Please sign in to comment.